feat(kafka): Implement basic async kafka consumer

Author: john-z-yang

.gitignore

@@ -2,3 +2,6 @@
 
 # Editors
 .DS_Store
+
+# Sqlite artifacts
+*.sqlite

Cargo.lock

@@ -6,14 +6,20 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-sentry_protos = "0.1.33"
+sentry_protos = "0.1.34"
+prost = "0.13"
+prost-types = "0.13.3"
 anyhow = "1.0.92"
 chrono = { version = "0.4.26" }
-sqlx = { version = "0.8.2", features = ["sqlite", "runtime-tokio", "chrono"] }
-prost = "0.13"
 tokio = { version = "1.41.0", features = ["full"] }
-prost-types = "0.13.3"
+tokio-util = "0.7.12"
+tokio-stream = { version = "0.1.16", features = ["full"] }
+async-stream = "0.3.5"
+futures = "0.3.31"
+tracing = "0.1.40"
+tracing-subscriber = "0.3.18"
 rdkafka = { version = "0.36.2", features = ["cmake-build"] }
+sqlx = { version = "0.8.2", features = ["sqlite", "runtime-tokio", "chrono"] }
 
 [dev-dependencies]
 rand = "0.8.5"

Cargo.toml

@@ -1,6 +1,7 @@
 CREATE TABLE IF NOT EXISTS inflight_taskactivations (
     id UUID NOT NULL PRIMARY KEY,
     activation BLOB NOT NULL,
+    partition INTEGER NOT NULL,
     offset BIGINTEGER NOT NULL,
     added_at DATETIME NOT NULL,
     deadletter_at DATETIME,

migrations/0001_create_inflight_taskactivations.sql

@@ -0,0 +1,25 @@
+use std::sync::Arc;
+
+use anyhow::{anyhow, Error};
+use chrono::Utc;
+use prost::Message as _;
+use rdkafka::{message::OwnedMessage, Message};
+use sentry_protos::sentry::v1::TaskActivation;
+
+use crate::inflight_activation_store::{InflightActivation, TaskActivationStatus};
+
+pub async fn deserialize_activation(msg: Arc<OwnedMessage>) -> Result<InflightActivation, Error> {
+    let Some(payload) = msg.payload() else {
+        return Err(anyhow!("Message has no payload"));
+    };
+    let activation = TaskActivation::decode(payload)?;
+    Ok(InflightActivation {
+        activation,
+        status: TaskActivationStatus::Pending,
+        partition: msg.partition(),
+        offset: msg.offset(),
+        added_at: Utc::now(),
+        deadletter_at: None,
+        processing_deadline: None,
+    })
+}

src/consumer/deserialize_activation.rs

@@ -0,0 +1,78 @@
+use std::{mem::replace, sync::Arc, time::Duration};
+
+use tracing::info;
+
+use crate::inflight_activation_store::{InflightActivation, InflightActivationStore};
+
+use super::kafka::{
+    ReduceConfig, ReduceShutdownBehaviour, ReduceShutdownCondition, Reducer,
+    ReducerWhenFullBehaviour,
+};
+
+pub struct InflightTaskWriterConfig {
+    pub max_buf_len: usize,
+    pub flush_interval: Option<Duration>,
+    pub when_full_behaviour: ReducerWhenFullBehaviour,
+    pub shutdown_behaviour: ReduceShutdownBehaviour,
+}
+
+pub struct InflightTaskWriter {
+    store: Arc<InflightActivationStore>,
+    buffer: Vec<InflightActivation>,
+    max_buf_len: usize,
+    reduce_config: ReduceConfig,
+}
+
+impl InflightTaskWriter {
+    pub fn new(store: Arc<InflightActivationStore>, config: InflightTaskWriterConfig) -> Self {
+        Self {
+            store,
+            buffer: Vec::with_capacity(config.max_buf_len),
+            max_buf_len: config.max_buf_len,
+            reduce_config: ReduceConfig {
+                shutdown_condition: ReduceShutdownCondition::Signal,
+                shutdown_behaviour: ReduceShutdownBehaviour::Flush,
+                when_full_behaviour: config.when_full_behaviour,
+                flush_interval: config.flush_interval,
+            },
+        }
+    }
+}
+
+impl Reducer for InflightTaskWriter {
+    type Input = InflightActivation;
+
+    type Output = ();
+
+    async fn reduce(&mut self, t: Self::Input) -> Result<(), anyhow::Error> {
+        self.buffer.push(t);
+        Ok(())
+    }
+
+    async fn flush(&mut self) -> Result<Self::Output, anyhow::Error> {
+        if self.buffer.is_empty() {
+            return Ok(());
+        }
+        let res = self
+            .store
+            .store(replace(
+                &mut self.buffer,
+                Vec::with_capacity(self.max_buf_len),
+            ))
+            .await?;
+        info!("Inserted {:?} entries", res.rows_affected);
+        Ok(())
+    }
+
+    fn reset(&mut self) {
+        self.buffer.clear();
+    }
+
+    fn is_full(&self) -> bool {
+        self.buffer.len() >= self.max_buf_len
+    }
+
+    fn get_reduce_config(&self) -> ReduceConfig {
+        self.reduce_config.clone()
+    }
+}