chore: add kafka producer to context (#25899)

PostHog · Oct 30, 2024 · 3ed93f7 · 3ed93f7
1 parent 019d4dc
commit 3ed93f7
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 7 deletions.
diff --git a/rust/cymbal/src/app_context.rs b/rust/cymbal/src/app_context.rs
@@ -1,8 +1,11 @@
-use std::time::Duration;
-
-use common_kafka::kafka_consumer::SingleTopicConsumer;
+use common_kafka::{
+    kafka_consumer::SingleTopicConsumer, kafka_producer::create_kafka_producer,
+    kafka_producer::KafkaContext,
+};
 use health::{HealthHandle, HealthRegistry};
+use rdkafka::producer::FutureProducer;
 use sqlx::{postgres::PgPoolOptions, PgPool};
+use std::time::Duration;
 use tracing::info;
 
 use crate::{
@@ -14,7 +17,8 @@ use crate::{
 pub struct AppContext {
     pub health_registry: HealthRegistry,
     pub worker_liveness: HealthHandle,
-    pub consumer: SingleTopicConsumer,
+    pub kafka_consumer: SingleTopicConsumer,
+    pub kafka_producer: FutureProducer<KafkaContext>,
     pub pool: PgPool,
     pub catalog: Catalog,
 }
@@ -25,8 +29,15 @@ impl AppContext {
         let worker_liveness = health_registry
             .register("worker".to_string(), Duration::from_secs(60))
             .await;
+        let kafka_liveness = health_registry
+            .register("rdkafka".to_string(), Duration::from_secs(30))
+            .await;
 
-        let consumer = SingleTopicConsumer::new(config.kafka.clone(), config.consumer.clone())?;
+        let kafka_consumer =
+            SingleTopicConsumer::new(config.kafka.clone(), config.consumer.clone())?;
+        let kafka_producer = create_kafka_producer(&config.kafka, kafka_liveness)
+            .await
+            .expect("failed to create kafka producer");
 
         let options = PgPoolOptions::new().max_connections(config.max_pg_connections);
         let pool = options.connect(&config.database_url).await?;
@@ -44,7 +55,8 @@ impl AppContext {
         Ok(Self {
             health_registry,
             worker_liveness,
-            consumer,
+            kafka_consumer,
+            kafka_producer,
             pool,
             catalog,
         })

diff --git a/rust/cymbal/src/main.rs b/rust/cymbal/src/main.rs
@@ -64,7 +64,7 @@ async fn main() -> Result<(), Error> {
         context.worker_liveness.report_healthy().await;
         // Just grab the event as a serde_json::Value and immediately drop it,
         // we can work out a real type for it later (once we're deployed etc)
-        let (event, offset): (ClickHouseEvent, _) = match context.consumer.json_recv().await {
+        let (event, offset): (ClickHouseEvent, _) = match context.kafka_consumer.json_recv().await {
             Ok(r) => r,
             Err(RecvErr::Kafka(e)) => {
                 return Err(e.into()); // Just die if we recieve a Kafka error