use std::collections::HashMap;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct StreamMessage {
pub offset: StreamOffset,
pub key: Option<Vec<u8>>,
pub payload: Vec<u8>,
pub timestamp: Option<u64>,
pub headers: HashMap<String, String>,
pub source: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum StreamOffset {
Kafka {
topic: String,
partition: i32,
offset: i64,
},
Redis {
stream: String,
id: String,
},
Nats {
subject: String,
sequence: u64,
},
Sequence(u64),
}
#[derive(Debug, Clone)]
pub struct StreamConsumerConfig {
pub group_id: String,
pub max_buffer_size: usize,
pub poll_interval: Duration,
pub commit_strategy: CommitStrategy,
pub max_poll_messages: usize,
pub connect_timeout: Duration,
}
impl Default for StreamConsumerConfig {
fn default() -> Self {
Self {
group_id: "ringkernel-default".to_string(),
max_buffer_size: 10_000,
poll_interval: Duration::from_millis(100),
commit_strategy: CommitStrategy::AfterProcess,
max_poll_messages: 500,
connect_timeout: Duration::from_secs(10),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommitStrategy {
AfterProcess,
AfterPoll,
Manual,
}
pub trait StreamConsumer: Send + Sync {
fn connect(&mut self) -> Result<(), StreamError>;
fn subscribe(&mut self, sources: &[&str]) -> Result<(), StreamError>;
fn poll(&mut self, max: usize) -> Result<Vec<StreamMessage>, StreamError>;
fn commit(&mut self, offset: &StreamOffset) -> Result<(), StreamError>;
fn disconnect(&mut self) -> Result<(), StreamError>;
fn is_connected(&self) -> bool;
fn metrics(&self) -> StreamConsumerMetrics;
}
#[derive(Debug, Clone, Default)]
pub struct StreamConsumerMetrics {
pub messages_consumed: u64,
pub messages_committed: u64,
pub errors: u64,
pub lag: u64,
pub throughput: f64,
pub connected_since: Option<Instant>,
}
#[derive(Debug, Clone)]
pub enum StreamError {
ConnectionFailed(String),
SubscribeFailed(String),
PollFailed(String),
CommitFailed(String),
Timeout(Duration),
ProviderError(String),
}
impl std::fmt::Display for StreamError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ConnectionFailed(msg) => write!(f, "Connection failed: {}", msg),
Self::SubscribeFailed(msg) => write!(f, "Subscribe failed: {}", msg),
Self::PollFailed(msg) => write!(f, "Poll failed: {}", msg),
Self::CommitFailed(msg) => write!(f, "Commit failed: {}", msg),
Self::Timeout(d) => write!(f, "Timeout after {:?}", d),
Self::ProviderError(msg) => write!(f, "Provider error: {}", msg),
}
}
}
impl std::error::Error for StreamError {}
pub struct InMemoryConsumer {
messages: Vec<StreamMessage>,
position: usize,
committed: Vec<StreamOffset>,
connected: bool,
subscriptions: Vec<String>,
metrics: StreamConsumerMetrics,
}
impl InMemoryConsumer {
pub fn new(messages: Vec<StreamMessage>) -> Self {
Self {
messages,
position: 0,
committed: Vec::new(),
connected: false,
subscriptions: Vec::new(),
metrics: StreamConsumerMetrics::default(),
}
}
pub fn empty() -> Self {
Self::new(Vec::new())
}
pub fn produce(&mut self, source: &str, key: Option<&[u8]>, payload: &[u8]) {
let seq = self.messages.len() as u64;
self.messages.push(StreamMessage {
offset: StreamOffset::Sequence(seq),
key: key.map(|k| k.to_vec()),
payload: payload.to_vec(),
timestamp: Some(seq),
headers: HashMap::new(),
source: source.to_string(),
});
}
pub fn committed_offsets(&self) -> &[StreamOffset] {
&self.committed
}
}
impl StreamConsumer for InMemoryConsumer {
fn connect(&mut self) -> Result<(), StreamError> {
self.connected = true;
self.metrics.connected_since = Some(Instant::now());
Ok(())
}
fn subscribe(&mut self, sources: &[&str]) -> Result<(), StreamError> {
self.subscriptions = sources.iter().map(|s| s.to_string()).collect();
Ok(())
}
fn poll(&mut self, max: usize) -> Result<Vec<StreamMessage>, StreamError> {
if !self.connected {
return Err(StreamError::ConnectionFailed("Not connected".into()));
}
let end = (self.position + max).min(self.messages.len());
let batch: Vec<StreamMessage> = self.messages[self.position..end]
.iter()
.filter(|m| self.subscriptions.is_empty() || self.subscriptions.contains(&m.source))
.cloned()
.collect();
self.position = end;
self.metrics.messages_consumed += batch.len() as u64;
Ok(batch)
}
fn commit(&mut self, offset: &StreamOffset) -> Result<(), StreamError> {
self.committed.push(offset.clone());
self.metrics.messages_committed += 1;
Ok(())
}
fn disconnect(&mut self) -> Result<(), StreamError> {
self.connected = false;
Ok(())
}
fn is_connected(&self) -> bool {
self.connected
}
fn metrics(&self) -> StreamConsumerMetrics {
self.metrics.clone()
}
}
#[derive(Debug, Clone)]
pub struct KafkaConfig {
pub brokers: String,
pub group_id: String,
pub auto_offset_reset: String,
pub enable_auto_commit: bool,
}
impl Default for KafkaConfig {
fn default() -> Self {
Self {
brokers: "localhost:9092".to_string(),
group_id: "ringkernel".to_string(),
auto_offset_reset: "earliest".to_string(),
enable_auto_commit: false,
}
}
}
#[derive(Debug, Clone)]
pub struct NatsConfig {
pub url: String,
pub credentials: Option<String>,
}
impl Default for NatsConfig {
fn default() -> Self {
Self {
url: "nats://localhost:4222".to_string(),
credentials: None,
}
}
}
#[derive(Debug, Clone)]
pub struct RedisStreamConfig {
pub url: String,
pub group: String,
pub consumer: String,
}
impl Default for RedisStreamConfig {
fn default() -> Self {
Self {
url: "redis://localhost:6379".to_string(),
group: "ringkernel".to_string(),
consumer: "worker-1".to_string(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_in_memory_consumer_basic() {
let mut consumer = InMemoryConsumer::empty();
consumer.produce("events", None, b"hello");
consumer.produce("events", None, b"world");
consumer.connect().unwrap();
consumer.subscribe(&["events"]).unwrap();
let msgs = consumer.poll(10).unwrap();
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[0].payload, b"hello");
assert_eq!(msgs[1].payload, b"world");
}
#[test]
fn test_in_memory_consumer_batching() {
let mut consumer = InMemoryConsumer::empty();
for i in 0..10 {
consumer.produce("topic", None, format!("msg-{}", i).as_bytes());
}
consumer.connect().unwrap();
consumer.subscribe(&["topic"]).unwrap();
let batch1 = consumer.poll(3).unwrap();
assert_eq!(batch1.len(), 3);
let batch2 = consumer.poll(3).unwrap();
assert_eq!(batch2.len(), 3);
let batch3 = consumer.poll(100).unwrap();
assert_eq!(batch3.len(), 4); }
#[test]
fn test_in_memory_consumer_commit() {
let mut consumer = InMemoryConsumer::empty();
consumer.produce("t", None, b"data");
consumer.connect().unwrap();
consumer.subscribe(&["t"]).unwrap();
let msgs = consumer.poll(1).unwrap();
consumer.commit(&msgs[0].offset).unwrap();
assert_eq!(consumer.committed_offsets().len(), 1);
}
#[test]
fn test_in_memory_consumer_not_connected() {
let mut consumer = InMemoryConsumer::empty();
let result = consumer.poll(1);
assert!(result.is_err());
}
#[test]
fn test_stream_consumer_metrics() {
let mut consumer = InMemoryConsumer::empty();
consumer.produce("t", None, b"a");
consumer.produce("t", None, b"b");
consumer.connect().unwrap();
consumer.subscribe(&["t"]).unwrap();
consumer.poll(10).unwrap();
let m = consumer.metrics();
assert_eq!(m.messages_consumed, 2);
assert!(m.connected_since.is_some());
}
#[test]
fn test_kafka_config_defaults() {
let config = KafkaConfig::default();
assert_eq!(config.brokers, "localhost:9092");
assert!(!config.enable_auto_commit);
}
#[test]
fn test_nats_config_defaults() {
let config = NatsConfig::default();
assert_eq!(config.url, "nats://localhost:4222");
}
#[test]
fn test_redis_config_defaults() {
let config = RedisStreamConfig::default();
assert_eq!(config.group, "ringkernel");
}
#[test]
fn test_stream_offset_equality() {
let a = StreamOffset::Sequence(42);
let b = StreamOffset::Sequence(42);
assert_eq!(a, b);
}
}