1use std::collections::HashMap;
20use std::time::{Duration, Instant};
21
22#[derive(Debug, Clone)]
24pub struct StreamMessage {
25 pub offset: StreamOffset,
27 pub key: Option<Vec<u8>>,
29 pub payload: Vec<u8>,
31 pub timestamp: Option<u64>,
33 pub headers: HashMap<String, String>,
35 pub source: String,
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
41pub enum StreamOffset {
42 Kafka {
44 topic: String,
46 partition: i32,
48 offset: i64,
50 },
51 Redis {
53 stream: String,
55 id: String,
57 },
58 Nats {
60 subject: String,
62 sequence: u64,
64 },
65 Sequence(u64),
67}
68
69#[derive(Debug, Clone)]
71pub struct StreamConsumerConfig {
72 pub group_id: String,
74 pub max_buffer_size: usize,
76 pub poll_interval: Duration,
78 pub commit_strategy: CommitStrategy,
80 pub max_poll_messages: usize,
82 pub connect_timeout: Duration,
84}
85
86impl Default for StreamConsumerConfig {
87 fn default() -> Self {
88 Self {
89 group_id: "ringkernel-default".to_string(),
90 max_buffer_size: 10_000,
91 poll_interval: Duration::from_millis(100),
92 commit_strategy: CommitStrategy::AfterProcess,
93 max_poll_messages: 500,
94 connect_timeout: Duration::from_secs(10),
95 }
96 }
97}
98
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub enum CommitStrategy {
102 AfterProcess,
104 AfterPoll,
106 Manual,
108}
109
110pub trait StreamConsumer: Send + Sync {
115 fn connect(&mut self) -> Result<(), StreamError>;
117
118 fn subscribe(&mut self, sources: &[&str]) -> Result<(), StreamError>;
120
121 fn poll(&mut self, max: usize) -> Result<Vec<StreamMessage>, StreamError>;
125
126 fn commit(&mut self, offset: &StreamOffset) -> Result<(), StreamError>;
128
129 fn disconnect(&mut self) -> Result<(), StreamError>;
131
132 fn is_connected(&self) -> bool;
134
135 fn metrics(&self) -> StreamConsumerMetrics;
137}
138
139#[derive(Debug, Clone, Default)]
141pub struct StreamConsumerMetrics {
142 pub messages_consumed: u64,
144 pub messages_committed: u64,
146 pub errors: u64,
148 pub lag: u64,
150 pub throughput: f64,
152 pub connected_since: Option<Instant>,
154}
155
156#[derive(Debug, Clone)]
158pub enum StreamError {
159 ConnectionFailed(String),
161 SubscribeFailed(String),
163 PollFailed(String),
165 CommitFailed(String),
167 Timeout(Duration),
169 ProviderError(String),
171}
172
173impl std::fmt::Display for StreamError {
174 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175 match self {
176 Self::ConnectionFailed(msg) => write!(f, "Connection failed: {}", msg),
177 Self::SubscribeFailed(msg) => write!(f, "Subscribe failed: {}", msg),
178 Self::PollFailed(msg) => write!(f, "Poll failed: {}", msg),
179 Self::CommitFailed(msg) => write!(f, "Commit failed: {}", msg),
180 Self::Timeout(d) => write!(f, "Timeout after {:?}", d),
181 Self::ProviderError(msg) => write!(f, "Provider error: {}", msg),
182 }
183 }
184}
185
186impl std::error::Error for StreamError {}
187
188pub struct InMemoryConsumer {
194 messages: Vec<StreamMessage>,
195 position: usize,
196 committed: Vec<StreamOffset>,
197 connected: bool,
198 subscriptions: Vec<String>,
199 metrics: StreamConsumerMetrics,
200}
201
202impl InMemoryConsumer {
203 pub fn new(messages: Vec<StreamMessage>) -> Self {
205 Self {
206 messages,
207 position: 0,
208 committed: Vec::new(),
209 connected: false,
210 subscriptions: Vec::new(),
211 metrics: StreamConsumerMetrics::default(),
212 }
213 }
214
215 pub fn empty() -> Self {
217 Self::new(Vec::new())
218 }
219
220 pub fn produce(&mut self, source: &str, key: Option<&[u8]>, payload: &[u8]) {
222 let seq = self.messages.len() as u64;
223 self.messages.push(StreamMessage {
224 offset: StreamOffset::Sequence(seq),
225 key: key.map(|k| k.to_vec()),
226 payload: payload.to_vec(),
227 timestamp: Some(seq),
228 headers: HashMap::new(),
229 source: source.to_string(),
230 });
231 }
232
233 pub fn committed_offsets(&self) -> &[StreamOffset] {
235 &self.committed
236 }
237}
238
239impl StreamConsumer for InMemoryConsumer {
240 fn connect(&mut self) -> Result<(), StreamError> {
241 self.connected = true;
242 self.metrics.connected_since = Some(Instant::now());
243 Ok(())
244 }
245
246 fn subscribe(&mut self, sources: &[&str]) -> Result<(), StreamError> {
247 self.subscriptions = sources.iter().map(|s| s.to_string()).collect();
248 Ok(())
249 }
250
251 fn poll(&mut self, max: usize) -> Result<Vec<StreamMessage>, StreamError> {
252 if !self.connected {
253 return Err(StreamError::ConnectionFailed("Not connected".into()));
254 }
255
256 let end = (self.position + max).min(self.messages.len());
257 let batch: Vec<StreamMessage> = self.messages[self.position..end]
258 .iter()
259 .filter(|m| self.subscriptions.is_empty() || self.subscriptions.contains(&m.source))
260 .cloned()
261 .collect();
262
263 self.position = end;
264 self.metrics.messages_consumed += batch.len() as u64;
265 Ok(batch)
266 }
267
268 fn commit(&mut self, offset: &StreamOffset) -> Result<(), StreamError> {
269 self.committed.push(offset.clone());
270 self.metrics.messages_committed += 1;
271 Ok(())
272 }
273
274 fn disconnect(&mut self) -> Result<(), StreamError> {
275 self.connected = false;
276 Ok(())
277 }
278
279 fn is_connected(&self) -> bool {
280 self.connected
281 }
282
283 fn metrics(&self) -> StreamConsumerMetrics {
284 self.metrics.clone()
285 }
286}
287
288#[derive(Debug, Clone)]
294pub struct KafkaConfig {
295 pub brokers: String,
297 pub group_id: String,
299 pub auto_offset_reset: String,
301 pub enable_auto_commit: bool,
303}
304
305impl Default for KafkaConfig {
306 fn default() -> Self {
307 Self {
308 brokers: "localhost:9092".to_string(),
309 group_id: "ringkernel".to_string(),
310 auto_offset_reset: "earliest".to_string(),
311 enable_auto_commit: false,
312 }
313 }
314}
315
316#[derive(Debug, Clone)]
318pub struct NatsConfig {
319 pub url: String,
321 pub credentials: Option<String>,
323}
324
325impl Default for NatsConfig {
326 fn default() -> Self {
327 Self {
328 url: "nats://localhost:4222".to_string(),
329 credentials: None,
330 }
331 }
332}
333
334#[derive(Debug, Clone)]
336pub struct RedisStreamConfig {
337 pub url: String,
339 pub group: String,
341 pub consumer: String,
343}
344
345impl Default for RedisStreamConfig {
346 fn default() -> Self {
347 Self {
348 url: "redis://localhost:6379".to_string(),
349 group: "ringkernel".to_string(),
350 consumer: "worker-1".to_string(),
351 }
352 }
353}
354
355#[cfg(test)]
356mod tests {
357 use super::*;
358
359 #[test]
360 fn test_in_memory_consumer_basic() {
361 let mut consumer = InMemoryConsumer::empty();
362 consumer.produce("events", None, b"hello");
363 consumer.produce("events", None, b"world");
364
365 consumer.connect().unwrap();
366 consumer.subscribe(&["events"]).unwrap();
367
368 let msgs = consumer.poll(10).unwrap();
369 assert_eq!(msgs.len(), 2);
370 assert_eq!(msgs[0].payload, b"hello");
371 assert_eq!(msgs[1].payload, b"world");
372 }
373
374 #[test]
375 fn test_in_memory_consumer_batching() {
376 let mut consumer = InMemoryConsumer::empty();
377 for i in 0..10 {
378 consumer.produce("topic", None, format!("msg-{}", i).as_bytes());
379 }
380
381 consumer.connect().unwrap();
382 consumer.subscribe(&["topic"]).unwrap();
383
384 let batch1 = consumer.poll(3).unwrap();
385 assert_eq!(batch1.len(), 3);
386
387 let batch2 = consumer.poll(3).unwrap();
388 assert_eq!(batch2.len(), 3);
389
390 let batch3 = consumer.poll(100).unwrap();
391 assert_eq!(batch3.len(), 4); }
393
394 #[test]
395 fn test_in_memory_consumer_commit() {
396 let mut consumer = InMemoryConsumer::empty();
397 consumer.produce("t", None, b"data");
398 consumer.connect().unwrap();
399 consumer.subscribe(&["t"]).unwrap();
400
401 let msgs = consumer.poll(1).unwrap();
402 consumer.commit(&msgs[0].offset).unwrap();
403
404 assert_eq!(consumer.committed_offsets().len(), 1);
405 }
406
407 #[test]
408 fn test_in_memory_consumer_not_connected() {
409 let mut consumer = InMemoryConsumer::empty();
410 let result = consumer.poll(1);
411 assert!(result.is_err());
412 }
413
414 #[test]
415 fn test_stream_consumer_metrics() {
416 let mut consumer = InMemoryConsumer::empty();
417 consumer.produce("t", None, b"a");
418 consumer.produce("t", None, b"b");
419
420 consumer.connect().unwrap();
421 consumer.subscribe(&["t"]).unwrap();
422 consumer.poll(10).unwrap();
423
424 let m = consumer.metrics();
425 assert_eq!(m.messages_consumed, 2);
426 assert!(m.connected_since.is_some());
427 }
428
429 #[test]
430 fn test_kafka_config_defaults() {
431 let config = KafkaConfig::default();
432 assert_eq!(config.brokers, "localhost:9092");
433 assert!(!config.enable_auto_commit);
434 }
435
436 #[test]
437 fn test_nats_config_defaults() {
438 let config = NatsConfig::default();
439 assert_eq!(config.url, "nats://localhost:4222");
440 }
441
442 #[test]
443 fn test_redis_config_defaults() {
444 let config = RedisStreamConfig::default();
445 assert_eq!(config.group, "ringkernel");
446 }
447
448 #[test]
449 fn test_stream_offset_equality() {
450 let a = StreamOffset::Sequence(42);
451 let b = StreamOffset::Sequence(42);
452 assert_eq!(a, b);
453 }
454}