# KafkaSource Implementation Summary
## Overview
The `KafkaSource` implementation provides a robust, production-ready Kafka consumer for ingesting events into the LLM Auto-Optimizer stream processing pipeline. It is built on top of `rdkafka` with comprehensive error handling, metrics tracking, and integration with the state backend for offset management.
## File Location
**Primary Implementation**: `/workspaces/llm-auto-optimizer/crates/processor/src/kafka/source.rs`
**Module Definition**: `/workspaces/llm-auto-optimizer/crates/processor/src/kafka/mod.rs`
**Examples**: `/workspaces/llm-auto-optimizer/crates/processor/examples/kafka_source.rs`
**Tests**: `/workspaces/llm-auto-optimizer/crates/processor/tests/kafka_source_integration.rs`
## Core Components
### 1. KafkaSource<T> Struct
The main consumer struct that manages Kafka consumption with the following fields:
- **consumer**: `Arc<StreamConsumer<SourceConsumerContext>>` - The underlying rdkafka consumer
- **config**: `KafkaSourceConfig` - Consumer configuration
- **metrics**: `Arc<RwLock<KafkaSourceMetrics>>` - Real-time metrics tracking
- **state_backend**: `Option<Arc<dyn StateBackend>>` - Optional state backend for offset persistence
- **shutdown**: `Arc<AtomicBool>` - Graceful shutdown signal
- **throughput_tracker**: `Arc<RwLock<ThroughputTracker>>` - Throughput calculation over time windows
### 2. Configuration (KafkaSourceConfig)
Comprehensive configuration with sensible defaults:
```rust
pub struct KafkaSourceConfig {
pub brokers: String, // Kafka brokers
pub group_id: String, // Consumer group ID
pub topics: Vec<String>, // Topics to subscribe
pub commit_strategy: CommitStrategy, // Offset commit strategy
pub commit_interval_secs: Option<u64>, // Periodic commit interval
pub auto_offset_reset: Option<String>, // earliest/latest
pub extra_config: HashMap<String, String>, // Additional Kafka config
pub max_retries: Option<u32>, // Max retries
pub retry_backoff_ms: Option<u64>, // Retry backoff
pub poll_timeout_ms: Option<u64>, // Poll timeout
}
```
### 3. Commit Strategies
Four commit strategies supported:
- **Auto**: Automatic commits by rdkafka (not recommended for exactly-once)
- **Manual**: Manual commit after each message
- **Periodic**: Automatic commits at specified intervals (recommended)
- **OnCheckpoint**: Commit only on checkpoint (for exactly-once with state backend)
### 4. Metrics Tracking (KafkaSourceMetrics)
Comprehensive metrics for monitoring:
```rust
pub struct KafkaSourceMetrics {
pub messages_consumed: u64, // Total messages consumed
pub messages_failed: u64, // Total failures
pub messages_per_partition: HashMap<i32, u64>, // Per-partition counts
pub lag_per_partition: HashMap<i32, i64>, // Estimated lag
pub bytes_consumed: u64, // Total bytes
pub throughput_msgs_per_sec: f64, // Messages/sec
pub throughput_bytes_per_sec: f64, // Bytes/sec
pub avg_message_size: f64, // Average message size
pub last_commit_time: Option<DateTime<Utc>>, // Last commit timestamp
pub total_commits: u64, // Total commits
pub rebalance_count: u64, // Rebalance events
pub assigned_partitions: Vec<i32>, // Current partitions
}
```
### 5. Message Envelope (KafkaMessage<T>)
Rich message structure with metadata:
```rust
pub struct KafkaMessage<T> {
pub payload: T, // Deserialized payload
pub key: Option<Vec<u8>>, // Message key
pub topic: String, // Topic name
pub partition: i32, // Partition number
pub offset: i64, // Message offset
pub timestamp: Option<DateTime<Utc>>, // Message timestamp
pub headers: HashMap<String, Vec<u8>>, // Message headers
}
```
### 6. Consumer Context (SourceConsumerContext)
Custom consumer context for handling rdkafka callbacks:
- **pre_rebalance**: Handles partition revocation
- **post_rebalance**: Handles partition assignment and updates metrics
- **commit_callback**: Tracks commit success/failure and updates metrics
### 7. Throughput Tracker
Calculates throughput over rolling time windows:
```rust
struct ThroughputTracker {
messages_in_window: u64,
bytes_in_window: u64,
window_start: Instant,
window_duration: Duration, // Default: 60 seconds
}
```
## Key Features Implemented
### 1. Async Message Consumption
- Built on `tokio` for non-blocking I/O
- Uses `StreamConsumer` from rdkafka
- Async/await throughout for efficient resource usage
### 2. Automatic Reconnection
- Handles transient errors with retry logic
- Configurable retry count and backoff
- Exponential backoff for connection errors
### 3. Consumer Group Management
- Full consumer group coordination via rdkafka
- Automatic partition rebalancing
- Pre/post rebalance callbacks for state management
### 4. Offset Management
- Multiple commit strategies (auto, manual, periodic, on-checkpoint)
- Integration with state backend for exactly-once processing
- Automatic offset commits based on strategy
- Commit error handling with logging
### 5. Partition Support
- Handles multiple partitions automatically
- Tracks metrics per partition
- Partition assignment/revocation callbacks
- Support for seeking to specific offsets
### 6. Graceful Shutdown
- Clean shutdown via atomic boolean flag
- Final commit on shutdown
- Proper resource cleanup
- Timeout handling for shutdown operations
### 7. Message Deserialization
- JSON deserialization using serde_json
- Type-safe deserialization to generic type T
- Error handling with detailed error messages
- Failed deserialization tracked in metrics
### 8. Metrics Tracking
- Real-time metrics updated on each message
- Throughput calculation over time windows
- Lag tracking per partition
- Commit and rebalance tracking
- Thread-safe metrics access via RwLock
## Methods
### Core Methods
#### `new(config, state_backend) -> Result<Self>`
Creates a new Kafka source consumer with the given configuration and optional state backend.
#### `subscribe(&self, topics) -> Result<()>`
Subscribes to the specified topics. Triggers partition assignment.
#### `poll(&self) -> Result<Option<KafkaMessage<T>>>`
Polls for a single message. Returns `None` on timeout or shutdown.
#### `commit(&self) -> Result<()>`
Commits current offsets to Kafka. Uses async commit mode.
#### `seek(&self, topic, partition, offset) -> Result<()>`
Seeks to a specific offset for a topic/partition.
#### `close(&self) -> Result<()>`
Gracefully stops the consumer with final commit.
### Advanced Methods
#### `start(&self, tx) -> Result<()>`
Starts the consumer loop, sending messages to a channel. Handles:
- Periodic commits based on strategy
- Automatic reconnection on errors
- Retry logic with backoff
- Graceful shutdown
#### `metrics(&self) -> KafkaSourceMetrics`
Returns a snapshot of current metrics.
#### `get_lag(&self) -> Result<HashMap<i32, i64>>`
Calculates consumer lag per partition by fetching watermarks from Kafka.
### Internal Methods
#### `create_consumer(config, context) -> Result<StreamConsumer>`
Creates and configures the underlying rdkafka consumer.
#### `process_message(&self, msg) -> Result<KafkaMessage<T>>`
Processes a borrowed message into a typed KafkaMessage with deserialization.
#### `update_metrics(&self, msg)`
Updates metrics after processing a message.
## Error Handling
The implementation handles various error scenarios:
1. **Connection Errors**: Automatic retry with backoff
2. **Deserialization Errors**: Logged and counted in metrics
3. **Commit Errors**: Logged with warnings (async commits)
4. **Partition EOF**: Handled gracefully, returns None
5. **Rebalance Errors**: Logged and tracked in metrics
All errors are properly typed using the `ProcessorError` enum from the error module.
## Thread Safety
- All shared state protected by `Arc` and `RwLock`
- Atomic operations for shutdown flag
- Safe concurrent access to metrics
- Thread-safe consumer from rdkafka
## Performance Optimizations
1. **Async I/O**: Non-blocking operations throughout
2. **Minimal Locking**: Read locks for metrics access, write locks only when updating
3. **Efficient Deserialization**: Direct from borrowed message bytes
4. **Batching**: rdkafka handles message batching internally
5. **Throughput Calculation**: Only calculated when window expires
6. **Channel Buffering**: Configurable buffer size for message channels
## Testing
### Unit Tests
Located at the bottom of `source.rs`:
- Configuration default values
- Commit strategy serialization
- Throughput tracker behavior
- Message structure validation
- Metrics default values
### Integration Tests
Located in `tests/kafka_source_integration.rs`:
- Consumer creation
- Topic subscription
- Message polling
- Metrics tracking
- Commit strategies
- Graceful shutdown
- Multiple topics
- Lag tracking
**Note**: Integration tests require a running Kafka instance and are marked with `#[ignore]`.
### Example Programs
Located in `examples/kafka_source.rs`:
- Periodic commit strategy example
- Manual commit strategy example
- Metrics tracking example
- Complete end-to-end usage
## Usage Examples
### Basic Usage
```rust
let config = KafkaSourceConfig {
brokers: "localhost:9092".to_string(),
group_id: "my-group".to_string(),
topics: vec!["events".to_string()],
commit_strategy: CommitStrategy::Periodic,
commit_interval_secs: Some(5),
..Default::default()
};
let source = KafkaSource::<MyEvent>::new(config, None).await?;
let (tx, mut rx) = mpsc::channel(1000);
tokio::spawn(async move {
source.start(tx).await.unwrap();
});
while let Some(msg) = rx.recv().await {
println!("Received: {:?}", msg.payload);
}
```
### With State Backend
```rust
let state_backend = Arc::new(MemoryStateBackend::new(None));
let config = KafkaSourceConfig {
commit_strategy: CommitStrategy::OnCheckpoint,
// ... other config
};
let source = KafkaSource::<MyEvent>::new(config, Some(state_backend)).await?;
```
### Manual Polling
```rust
let source = KafkaSource::<MyEvent>::new(config, None).await?;
source.subscribe(&["topic".to_string()]).await?;
loop {
match source.poll().await? {
Some(msg) => {
process(&msg.payload);
source.commit().await?;
}
None => break,
}
}
source.close().await?;
```
## Integration with Pipeline
The KafkaSource integrates with the stream processing pipeline:
1. **Event Ingestion**: Consumes events from Kafka topics
2. **Deserialization**: Converts to ProcessorEvent or custom types
3. **Windowing**: Events flow into windowing operators
4. **Aggregation**: Aggregated within windows
5. **Offset Management**: Coordinates with state backend for exactly-once
## Dependencies
- **rdkafka**: Kafka client library (v0.36+)
- **tokio**: Async runtime
- **serde/serde_json**: Serialization/deserialization
- **chrono**: Timestamp handling
- **tracing**: Logging and instrumentation
- **dashmap**: Not used (can be removed)
- **async_trait**: Not used (can be removed)
## Future Enhancements
Potential improvements:
1. **Custom Deserializers**: Support for Avro, Protobuf, etc.
2. **Pause/Resume**: Ability to pause/resume consumption
3. **Offset Reset**: Manual offset reset to specific timestamp/offset
4. **Multi-threaded Consumption**: Parallel processing of partitions
5. **Backpressure**: Better backpressure handling
6. **Schema Registry**: Integration with Confluent Schema Registry
7. **Dead Letter Queue**: Failed message routing
8. **Metrics Export**: Prometheus/OpenTelemetry integration
## Conclusion
The KafkaSource implementation provides a production-ready, feature-rich Kafka consumer that integrates seamlessly with the stream processing pipeline. It handles all aspects of Kafka consumption including offset management, partition rebalancing, error handling, and metrics tracking while maintaining high performance and reliability.
The implementation follows Rust best practices with proper error handling, thread safety, and comprehensive testing. It is ready for production use in the LLM Auto-Optimizer system.