feat(hydro_test): add Kafka example with producer/consumer module

akainth015 · akainth015 · commit 4983720633e8 · 2026-04-27T20:30:01.000Z
- Add hydro_test/src/kafka/mod.rs with kafka_producer, kafka_consumer, dest_kafka, and setup_topic helpers following the SQS PR #2746 pattern - Complete hydro_test/examples/kafka.rs: leader produces 1M financial transactions, consumer cluster computes per-account balances - Add 'kafka' feature flag gating rdkafka and futures-util as optional deps - Add llvm-tools component to rust-toolchain.toml for rust-lld
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/hydro_test/Cargo.toml b/hydro_test/Cargo.toml
@@ -12,6 +12,11 @@ docker = ["hydro_lang/docker_deploy"]
 ecs = ["hydro_lang/ecs_deploy"]
 maelstrom = ["hydro_lang/maelstrom"]
 stageleft_macro_entrypoint = ["hydro_lang/stageleft_macro_entrypoint"]
+kafka = ["dep:rdkafka", "dep:futures-util"]
+
+[[example]]
+name = "kafka"
+required-features = ["kafka"]
 
 [dependencies]
 hydro_lang = { path = "../hydro_lang", version = "^0.15.0" }
@@ -33,7 +38,8 @@ bytes = "1.1.0"
 # https://github.com/GitoxideLabs/cargo-smart-release/issues/36
 example_test = { path = "../example_test", version = "^0.0.0", optional = true }
 hydro_build_utils = { path = "../hydro_build_utils", version = "^0.0.1", optional = true }
-rdkafka = "0.39.0"
+rdkafka = { version = "0.39.0", optional = true }
+futures-util = { version = "0.3.0", default-features = false, features = ["alloc"], optional = true }
 
 [build-dependencies]
 stageleft_tool.workspace = true
diff --git a/hydro_test/examples/kafka.rs b/hydro_test/examples/kafka.rs
@@ -1,12 +1,181 @@
-fn main() {
-    // [Leader]
-    // If it already exists, delete the "financial_transactions" topic
-    // Create a new topic named "financial_transactions" with 10 partitions.
+use std::collections::HashMap;
+use std::sync::Arc;
 
-    // Produce 1 million transactions into the financial_transactions topic, spread evenly
-    // across the 10 partitions
+use clap::{ArgAction, Parser};
+use hydro_deploy::gcp::GcpNetwork;
+use hydro_deploy::{AwsNetwork, Deployment, Host};
+use hydro_lang::deploy::TrybuildHost;
+use hydro_lang::live_collections::stream::{ExactlyOnce, TotalOrder};
+use hydro_lang::location::Location;
+use hydro_lang::nondet::nondet;
+use hydro_lang::viz::config::GraphConfig;
+use hydro_test::kafka::{dest_kafka, kafka_consumer, kafka_producer, setup_topic};
+use stageleft::q;
 
-    // [Consumers]
-    // Once the cluster of consumers receives the go-ahead from the leader stream, it may then
-    // begin to consume the financial transactions and compute the balance of each bank account
+type HostCreator = Box<dyn Fn(&mut Deployment) -> Arc<dyn Host>>;
+
+const TOPIC: &str = "financial_transactions";
+const NUM_PARTITIONS: i32 = 10;
+const NUM_TRANSACTIONS: usize = 1_000_000;
+const NUM_CONSUMERS: usize = 3;
+
+// cargo run -p hydro_test --example kafka --features kafka -- --brokers 'localhost:9092'
+#[derive(Parser, Debug)]
+#[command(group(
+    clap::ArgGroup::new("cloud")
+        .args(&["gcp", "aws"])
+        .multiple(false)
+))]
+struct Args {
+    #[clap(flatten)]
+    graph: GraphConfig,
+
+    /// Use GCP for deployment (provide project name)
+    #[arg(long)]
+    gcp: Option<String>,
+
+    /// Use AWS, make sure credentials are set up
+    #[arg(long, action = ArgAction::SetTrue)]
+    aws: bool,
+
+    /// Kafka bootstrap servers
+    #[arg(long, default_value = "localhost:9092")]
+    brokers: String,
+}
+
+enum Leader {}
+enum Consumer {}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+    let mut deployment = Deployment::new();
+
+    let create_host: HostCreator = if let Some(project) = &args.gcp {
+        let network = GcpNetwork::new(project, None);
+        let project = project.clone();
+
+        Box::new(move |deployment| -> Arc<dyn Host> {
+            deployment
+                .GcpComputeEngineHost()
+                .project(&project)
+                .machine_type("e2-micro")
+                .image("debian-cloud/debian-11")
+                .region("us-west1-a")
+                .network(network.clone())
+                .add()
+        })
+    } else if args.aws {
+        let region = "us-east-1";
+        let network = AwsNetwork::new(region, None);
+
+        Box::new(move |deployment| -> Arc<dyn Host> {
+            deployment
+                .AwsEc2Host()
+                .region(region)
+                .instance_type("t3.micro")
+                .ami("ami-0e95a5e2743ec9ec9") // Amazon Linux 2
+                .network(network.clone())
+                .add()
+        })
+    } else {
+        let localhost = deployment.Localhost();
+        Box::new(move |_| -> Arc<dyn Host> { localhost.clone() })
+    };
+
+    // [Leader] Setup topic and produce transactions
+    setup_topic(&args.brokers, TOPIC, NUM_PARTITIONS).await;
+    println!("Topic '{}' created with {} partitions", TOPIC, NUM_PARTITIONS);
+
+    let mut flow = hydro_lang::compile::builder::FlowBuilder::new();
+    let leader = flow.process::<Leader>();
+    let consumers = flow.cluster::<Consumer>();
+
+    // Leader: produce 1M transactions spread across 10 partitions.
+    // Each transaction is (account_id, amount) serialized as key=account, value=amount.
+    {
+        let producer = kafka_producer(&leader, &args.brokers);
+        let transactions = leader.source_iter(q!({
+            (0..NUM_TRANSACTIONS).map(|i| {
+                let account = format!("account_{}", i % 100);
+                let amount = format!("{}", (i % 201) as i64 - 100); // range [-100, 100]
+                (account, amount)
+            })
+        }));
+        dest_kafka(producer, transactions, TOPIC);
+    }
+
+    // Consumers: read from topic and compute per-account balances.
+    {
+        let messages = kafka_consumer(&consumers, &args.brokers, "kafka_example_consumers", TOPIC)
+            .assume_ordering::<TotalOrder>(
+                nondet!(/** Safe: side effect is only printing final balances. */),
+            )
+            .assume_retries::<ExactlyOnce>(
+                nondet!(/** Safe: side effect is only printing final balances. */),
+            )
+            .filter_map(q!(|msg| {
+                let key = rdkafka::Message::key(&msg)
+                    .map(|k| String::from_utf8_lossy(k).to_string())?;
+                let value = rdkafka::Message::payload(&msg)
+                    .map(|v| String::from_utf8_lossy(v).to_string())?;
+                let amount: i64 = value.parse().ok()?;
+                Some((key, amount))
+            }))
+            .fold(
+                q!(|| HashMap::<String, i64>::new()),
+                q!(|balances, (account, amount)| {
+                    *balances.entry(account).or_insert(0) += amount;
+                }),
+            );
+
+        messages.into_stream().for_each(q!(|balances: HashMap<String, i64>| {
+            println!("Final balances ({} accounts):", balances.len());
+            let mut sorted: Vec<_> = balances.into_iter().collect();
+            sorted.sort_by(|a, b| a.0.cmp(&b.0));
+            for (account, balance) in sorted.iter().take(10) {
+                println!("  {}: {}", account, balance);
+            }
+            if sorted.len() > 10 {
+                println!("  ... and {} more accounts", sorted.len() - 10);
+            }
+        }));
+    }
+
+    // Extract the IR BEFORE the builder is consumed by deployment methods
+    let built = flow.finalize();
+
+    // Generate graph visualizations based on command line arguments
+    built.generate_graph_with_config(&args.graph, None)?;
+
+    // If we're just generating a graph file, exit early
+    if args.graph.should_exit_after_graph_generation() {
+        return Ok(());
+    }
+
+    // Now use the built flow for deployment with optimization
+    let _nodes = built
+        .with_default_optimize()
+        .with_process(
+            &leader,
+            TrybuildHost::new(create_host(&mut deployment))
+                .features(vec!["kafka".to_owned()]),
+        )
+        .with_cluster(
+            &consumers,
+            (0..NUM_CONSUMERS).map(|_| {
+                TrybuildHost::new(create_host(&mut deployment))
+                    .features(vec!["kafka".to_owned()])
+            }),
+        )
+        .deploy(&mut deployment);
+
+    deployment.deploy().await.unwrap();
+    deployment.start().await.unwrap();
+
+    println!("Running Kafka financial transactions example...");
+    println!("Press Ctrl+C to stop.");
+
+    tokio::signal::ctrl_c().await.unwrap();
+    Ok(())
 }
diff --git a/hydro_test/src/kafka/mod.rs b/hydro_test/src/kafka/mod.rs
@@ -0,0 +1,152 @@
+use std::future::Future;
+
+use hydro_lang::live_collections::boundedness::Boundedness;
+use hydro_lang::live_collections::stream::{AtLeastOnce, ExactlyOnce, NoOrder, Ordering};
+use hydro_lang::location::tick::{NoAtomic, NoTick};
+use hydro_lang::location::Location;
+use hydro_lang::prelude::*;
+use rdkafka::message::OwnedMessage;
+use rdkafka::producer::FutureProducer;
+
+#[ctor::ctor]
+fn init_rewrites() {
+    stageleft::add_private_reexport(
+        vec!["rdkafka", "producer", "future_producer"],
+        vec!["rdkafka", "producer"],
+    );
+    stageleft::add_private_reexport(
+        vec!["rdkafka", "consumer", "stream_consumer"],
+        vec!["rdkafka", "consumer"],
+    );
+    stageleft::add_private_reexport(
+        vec!["rdkafka", "message", "owned_message"],
+        vec!["rdkafka", "message"],
+    );
+    stageleft::add_private_reexport(
+        vec!["futures_util", "stream", "stream"],
+        vec!["futures_util", "stream"],
+    );
+    stageleft::add_private_reexport(
+        vec!["futures_util", "stream", "unfold"],
+        vec!["futures_util", "stream"],
+    );
+}
+
+/// Creates a Kafka `FutureProducer` singleton.
+pub fn kafka_producer<'a, Loc>(
+    location: &Loc,
+    brokers: &'a str,
+) -> Singleton<FutureProducer, Loc, Bounded>
+where
+    Loc: Location<'a> + NoTick + NoAtomic,
+{
+    location.singleton(q!({
+        rdkafka::config::ClientConfig::new()
+            .set("bootstrap.servers", brokers)
+            .create::<rdkafka::producer::FutureProducer>()
+            .expect("Failed to create Kafka producer")
+    }))
+}
+
+/// Consumes messages from a Kafka topic. Returns at-least-once, unordered delivery.
+pub fn kafka_consumer<'a, Loc>(
+    location: &Loc,
+    brokers: &'a str,
+    group_id: &'a str,
+    topic: &'a str,
+) -> Stream<OwnedMessage, Loc, Bounded, NoOrder, AtLeastOnce>
+where
+    Loc: Location<'a> + NoTick + NoAtomic,
+{
+    location
+        .singleton(q!({
+            let consumer: rdkafka::consumer::StreamConsumer =
+                rdkafka::config::ClientConfig::new()
+                    .set("bootstrap.servers", brokers)
+                    .set("group.id", group_id)
+                    .set("auto.offset.reset", "earliest")
+                    .create()
+                    .expect("Failed to create Kafka consumer");
+            rdkafka::consumer::Consumer::subscribe(&consumer, &[topic])
+                .expect("Failed to subscribe to topic");
+            std::sync::Arc::new(consumer)
+        }))
+        .into_stream()
+        .flat_map_stream_blocking(q!(
+            |consumer: std::sync::Arc<rdkafka::consumer::StreamConsumer>| {
+                futures_util::stream::unfold(consumer, |consumer| async move {
+                    loop {
+                        match rdkafka::consumer::StreamConsumer::recv(&*consumer).await {
+                            Ok(msg) => {
+                                return Some((rdkafka::message::BorrowedMessage::detach(&msg), consumer));
+                            }
+                            Err(e) => {
+                                eprintln!("Kafka consumer error: {}", e);
+                                continue;
+                            }
+                        }
+                    }
+                })
+            }
+        ))
+        .weaken_retries()
+        .weaken_ordering()
+}
+
+/// Sends `(key, payload)` pairs to a Kafka topic.
+pub fn dest_kafka<'a, Loc, Bound: Boundedness, Order: Ordering>(
+    producer: Singleton<FutureProducer, Loc, Bounded>,
+    input: Stream<(String, String), Loc, Bound, Order, ExactlyOnce>,
+    topic: &'a str,
+) where
+    Loc: Location<'a>,
+{
+    input
+        .cross_singleton(producer)
+        .map(q!(
+            |((key, payload), producer)| self::kafka_send(producer, topic, key, payload)
+        ))
+        .resolve_futures_blocking();
+}
+
+fn kafka_send(
+    producer: FutureProducer,
+    topic: &str,
+    key: String,
+    payload: String,
+) -> impl Future<Output = ()> {
+    let topic = topic.to_owned();
+    async move {
+        let record = rdkafka::producer::FutureRecord::to(&topic)
+            .key(&key)
+            .payload(&payload);
+        producer
+            .send(record, rdkafka::util::Timeout::Never)
+            .await
+            .expect("Failed to send message to Kafka");
+    }
+}
+
+/// Admin helper: delete topic if it exists, then create it with the given number of partitions.
+pub async fn setup_topic(brokers: &str, topic: &str, num_partitions: i32) {
+    use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication};
+    use rdkafka::config::ClientConfig;
+
+    let admin: AdminClient<rdkafka::client::DefaultClientContext> = ClientConfig::new()
+        .set("bootstrap.servers", brokers)
+        .create()
+        .expect("Failed to create Kafka admin client");
+
+    let opts = AdminOptions::new();
+
+    // Delete topic if it exists (ignore errors if it doesn't exist)
+    let _ = admin.delete_topics(&[topic], &opts).await;
+    // Brief pause to let deletion propagate
+    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+
+    let new_topic = NewTopic::new(topic, num_partitions, TopicReplication::Fixed(1));
+    admin
+        .create_topics(&[new_topic], &opts)
+        .await
+        .expect("Failed to create Kafka topic");
+}
diff --git a/hydro_test/src/lib.rs b/hydro_test/src/lib.rs
@@ -9,6 +9,9 @@ pub mod local;
 pub mod maelstrom;
 pub mod tutorials;
 
+#[cfg(feature = "kafka")]
+pub mod kafka;
+
 #[doc(hidden)]
 #[cfg(doctest)]
 mod docs {
diff --git a/rust-toolchain.toml b/rust-toolchain.toml
@@ -5,5 +5,6 @@ components = [
     "clippy",
     # https://github.com/dtolnay/trybuild?tab=readme-ov-file#troubleshooting
     "rust-src",
+    "llvm-tools",
 ]
 targets = ["wasm32-unknown-unknown", "x86_64-unknown-linux-musl"]

Original file line number	Diff line number	Diff line change
`@@ -5,5 +5,6 @@ components = [`
`5`	`5`	`"clippy",`
`6`	`6`	`# https://github.com/dtolnay/trybuild?tab=readme-ov-file#troubleshooting`
`7`	`7`	`"rust-src",`
	`8`	`+ "llvm-tools",`
`8`	`9`	`]`
`9`	`10`	`targets = ["wasm32-unknown-unknown", "x86_64-unknown-linux-musl"]`