llm_optimizer_processor/kafka/
offset_manager.rs

1//! Kafka offset management and checkpointing
2//!
3//! This module provides comprehensive offset tracking, persistence, and recovery
4//! capabilities for Kafka consumers. It integrates with the state backend system
5//! to ensure exactly-once processing semantics and fault tolerance.
6//!
7//! # Features
8//!
9//! - **Per-partition offset tracking**: Track offsets independently for each partition
10//! - **Watermark-based commits**: Commit offsets only after watermark passes
11//! - **Checkpoint integration**: Save offsets as part of state checkpoints
12//! - **Offset lag monitoring**: Calculate and track consumer lag
13//! - **Reset strategies**: Support for earliest, latest, and specific offset resets
14//! - **Automatic commit strategies**: Periodic, watermark-based, or manual commits
15//!
16//! # Example
17//!
18//! ```rust,no_run
19//! use processor::kafka::{OffsetManager, OffsetCommitStrategy, InMemoryOffsetStore};
20//! use processor::state::MemoryStateBackend;
21//! use std::sync::Arc;
22//! use std::time::Duration;
23//!
24//! #[tokio::main]
25//! async fn main() -> anyhow::Result<()> {
26//!     let backend = Arc::new(MemoryStateBackend::new(None));
27//!     let store = Arc::new(InMemoryOffsetStore::new());
28//!
29//!     let mut manager = OffsetManager::new(
30//!         "my-consumer-group".to_string(),
31//!         backend,
32//!         store,
33//!         OffsetCommitStrategy::Periodic {
34//!             interval: Duration::from_secs(5),
35//!         },
36//!     );
37//!
38//!     // Update offset after processing event
39//!     manager.update_offset("my-topic", 0, 100).await?;
40//!
41//!     // Commit offsets
42//!     manager.commit().await?;
43//!
44//!     // Check lag
45//!     let lag = manager.lag("my-topic", 0, 200).await?;
46//!     println!("Current lag: {}", lag);
47//!
48//!     Ok(())
49//! }
50//! ```
51
52use async_trait::async_trait;
53use chrono::{DateTime, Utc};
54use dashmap::DashMap;
55use serde::{Deserialize, Serialize};
56use std::collections::HashMap;
57use std::sync::Arc;
58use std::time::Duration;
59use tokio::sync::{Mutex, RwLock};
60use tracing::{debug, error, info, warn};
61
62use crate::error::{StateError, StateResult};
63use crate::state::StateBackend;
64use crate::watermark::Watermark;
65
66/// Represents a topic-partition pair
67#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
68pub struct TopicPartition {
69    /// Kafka topic name
70    pub topic: String,
71    /// Partition number
72    pub partition: u32,
73}
74
75impl TopicPartition {
76    /// Create a new TopicPartition
77    pub fn new(topic: String, partition: u32) -> Self {
78        Self { topic, partition }
79    }
80
81    /// Convert to string key for storage
82    pub fn to_key(&self) -> String {
83        format!("{}:{}", self.topic, self.partition)
84    }
85}
86
87impl std::fmt::Display for TopicPartition {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        write!(f, "{}:{}", self.topic, self.partition)
90    }
91}
92
93/// Offset information for a partition
94#[derive(Debug, Clone, Serialize, Deserialize)]
95pub struct OffsetInfo {
96    /// Current committed offset
97    pub offset: i64,
98    /// High watermark (latest available offset in partition)
99    pub high_watermark: Option<i64>,
100    /// Timestamp when offset was last updated
101    pub last_updated: DateTime<Utc>,
102    /// Metadata associated with this offset
103    pub metadata: Option<String>,
104}
105
106impl OffsetInfo {
107    /// Create new offset info
108    pub fn new(offset: i64) -> Self {
109        Self {
110            offset,
111            high_watermark: None,
112            last_updated: Utc::now(),
113            metadata: None,
114        }
115    }
116
117    /// Create with high watermark
118    pub fn with_high_watermark(offset: i64, high_watermark: i64) -> Self {
119        Self {
120            offset,
121            high_watermark: Some(high_watermark),
122            last_updated: Utc::now(),
123            metadata: None,
124        }
125    }
126
127    /// Calculate lag (difference between high watermark and current offset)
128    pub fn lag(&self) -> Option<i64> {
129        self.high_watermark.map(|hwm| hwm - self.offset)
130    }
131}
132
133/// Strategy for committing offsets
134#[derive(Debug, Clone, Serialize, Deserialize)]
135pub enum OffsetCommitStrategy {
136    /// Commit periodically at fixed intervals
137    Periodic {
138        /// Interval between commits
139        interval: Duration,
140    },
141    /// Commit after processing N events
142    EveryNEvents {
143        /// Number of events between commits
144        n: usize,
145    },
146    /// Commit when watermark advances
147    OnWatermark,
148    /// Manual commit only (application controls)
149    Manual,
150}
151
152impl Default for OffsetCommitStrategy {
153    fn default() -> Self {
154        Self::Periodic {
155            interval: Duration::from_secs(5),
156        }
157    }
158}
159
160/// Reset strategy for partition offsets
161#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
162pub enum OffsetResetStrategy {
163    /// Reset to earliest available offset
164    Earliest,
165    /// Reset to latest available offset (end of partition)
166    Latest,
167    /// Reset to specific offset
168    Specific(i64),
169    /// Reset to specific timestamp
170    Timestamp(i64),
171}
172
173/// Statistics about offset management
174#[derive(Debug, Clone, Default, Serialize, Deserialize)]
175pub struct OffsetStats {
176    /// Total number of commits
177    pub total_commits: u64,
178    /// Number of successful commits
179    pub successful_commits: u64,
180    /// Number of failed commits
181    pub failed_commits: u64,
182    /// Total offsets committed
183    pub total_offsets_committed: u64,
184    /// Last commit timestamp
185    pub last_commit_time: Option<DateTime<Utc>>,
186    /// Average commit duration in milliseconds
187    pub avg_commit_duration_ms: u64,
188}
189
190/// Trait for offset storage backends
191///
192/// Implementations can store offsets in memory, state backends,
193/// external systems, or Kafka itself.
194#[async_trait]
195pub trait OffsetStore: Send + Sync {
196    /// Store offset for a topic-partition
197    ///
198    /// # Arguments
199    /// * `consumer_group` - Consumer group ID
200    /// * `topic_partition` - Topic and partition
201    /// * `offset_info` - Offset information to store
202    ///
203    /// # Returns
204    /// * `Ok(())` - Offset stored successfully
205    /// * `Err(StateError)` - Storage failed
206    async fn store_offset(
207        &self,
208        consumer_group: &str,
209        topic_partition: &TopicPartition,
210        offset_info: &OffsetInfo,
211    ) -> StateResult<()>;
212
213    /// Retrieve offset for a topic-partition
214    ///
215    /// # Arguments
216    /// * `consumer_group` - Consumer group ID
217    /// * `topic_partition` - Topic and partition
218    ///
219    /// # Returns
220    /// * `Ok(Some(OffsetInfo))` - Offset found
221    /// * `Ok(None)` - No offset stored
222    /// * `Err(StateError)` - Retrieval failed
223    async fn get_offset(
224        &self,
225        consumer_group: &str,
226        topic_partition: &TopicPartition,
227    ) -> StateResult<Option<OffsetInfo>>;
228
229    /// Get all offsets for a consumer group
230    ///
231    /// # Arguments
232    /// * `consumer_group` - Consumer group ID
233    ///
234    /// # Returns
235    /// * `Ok(HashMap)` - Map of topic-partitions to offset info
236    /// * `Err(StateError)` - Retrieval failed
237    async fn get_all_offsets(
238        &self,
239        consumer_group: &str,
240    ) -> StateResult<HashMap<TopicPartition, OffsetInfo>>;
241
242    /// Delete offset for a topic-partition
243    ///
244    /// # Arguments
245    /// * `consumer_group` - Consumer group ID
246    /// * `topic_partition` - Topic and partition
247    ///
248    /// # Returns
249    /// * `Ok(())` - Offset deleted
250    /// * `Err(StateError)` - Deletion failed
251    async fn delete_offset(
252        &self,
253        consumer_group: &str,
254        topic_partition: &TopicPartition,
255    ) -> StateResult<()>;
256
257    /// Clear all offsets for a consumer group
258    ///
259    /// # Arguments
260    /// * `consumer_group` - Consumer group ID
261    ///
262    /// # Returns
263    /// * `Ok(())` - All offsets cleared
264    /// * `Err(StateError)` - Clear failed
265    async fn clear_offsets(&self, consumer_group: &str) -> StateResult<()>;
266}
267
268/// In-memory offset store implementation
269///
270/// Stores offsets in memory using a concurrent hash map.
271/// Suitable for testing and development.
272pub struct InMemoryOffsetStore {
273    offsets: Arc<DashMap<String, HashMap<TopicPartition, OffsetInfo>>>,
274}
275
276impl InMemoryOffsetStore {
277    /// Create a new in-memory offset store
278    pub fn new() -> Self {
279        Self {
280            offsets: Arc::new(DashMap::new()),
281        }
282    }
283}
284
285impl Default for InMemoryOffsetStore {
286    fn default() -> Self {
287        Self::new()
288    }
289}
290
291#[async_trait]
292impl OffsetStore for InMemoryOffsetStore {
293    async fn store_offset(
294        &self,
295        consumer_group: &str,
296        topic_partition: &TopicPartition,
297        offset_info: &OffsetInfo,
298    ) -> StateResult<()> {
299        self.offsets
300            .entry(consumer_group.to_string())
301            .or_insert_with(HashMap::new)
302            .insert(topic_partition.clone(), offset_info.clone());
303        Ok(())
304    }
305
306    async fn get_offset(
307        &self,
308        consumer_group: &str,
309        topic_partition: &TopicPartition,
310    ) -> StateResult<Option<OffsetInfo>> {
311        Ok(self
312            .offsets
313            .get(consumer_group)
314            .and_then(|offsets| offsets.get(topic_partition).cloned()))
315    }
316
317    async fn get_all_offsets(
318        &self,
319        consumer_group: &str,
320    ) -> StateResult<HashMap<TopicPartition, OffsetInfo>> {
321        Ok(self
322            .offsets
323            .get(consumer_group)
324            .map(|offsets| offsets.clone())
325            .unwrap_or_default())
326    }
327
328    async fn delete_offset(
329        &self,
330        consumer_group: &str,
331        topic_partition: &TopicPartition,
332    ) -> StateResult<()> {
333        if let Some(mut offsets) = self.offsets.get_mut(consumer_group) {
334            offsets.remove(topic_partition);
335        }
336        Ok(())
337    }
338
339    async fn clear_offsets(&self, consumer_group: &str) -> StateResult<()> {
340        self.offsets.remove(consumer_group);
341        Ok(())
342    }
343}
344
345/// State backend offset store implementation
346///
347/// Stores offsets in a state backend for durability and recovery.
348pub struct StateBackendOffsetStore<B: StateBackend> {
349    backend: Arc<B>,
350    prefix: String,
351}
352
353impl<B: StateBackend> StateBackendOffsetStore<B> {
354    /// Create a new state backend offset store
355    ///
356    /// # Arguments
357    /// * `backend` - State backend to use
358    /// * `prefix` - Key prefix for offset storage (default: "kafka_offset:")
359    pub fn new(backend: Arc<B>) -> Self {
360        Self::with_prefix(backend, "kafka_offset:".to_string())
361    }
362
363    /// Create with custom prefix
364    pub fn with_prefix(backend: Arc<B>, prefix: String) -> Self {
365        Self { backend, prefix }
366    }
367
368    /// Build storage key
369    fn build_key(&self, consumer_group: &str, topic_partition: &TopicPartition) -> Vec<u8> {
370        format!(
371            "{}{}:{}",
372            self.prefix,
373            consumer_group,
374            topic_partition.to_key()
375        )
376        .into_bytes()
377    }
378
379    /// Build prefix for consumer group
380    fn build_group_prefix(&self, consumer_group: &str) -> Vec<u8> {
381        format!("{}{}:", self.prefix, consumer_group).into_bytes()
382    }
383}
384
385#[async_trait]
386impl<B: StateBackend> OffsetStore for StateBackendOffsetStore<B> {
387    async fn store_offset(
388        &self,
389        consumer_group: &str,
390        topic_partition: &TopicPartition,
391        offset_info: &OffsetInfo,
392    ) -> StateResult<()> {
393        let key = self.build_key(consumer_group, topic_partition);
394        let value = bincode::serialize(offset_info).map_err(|e| {
395            StateError::SerializationFailed {
396                key: String::from_utf8_lossy(&key).to_string(),
397                reason: e.to_string(),
398            }
399        })?;
400
401        self.backend.put(&key, &value).await
402    }
403
404    async fn get_offset(
405        &self,
406        consumer_group: &str,
407        topic_partition: &TopicPartition,
408    ) -> StateResult<Option<OffsetInfo>> {
409        let key = self.build_key(consumer_group, topic_partition);
410
411        match self.backend.get(&key).await? {
412            Some(value) => {
413                let offset_info = bincode::deserialize(&value).map_err(|e| {
414                    StateError::DeserializationFailed {
415                        key: String::from_utf8_lossy(&key).to_string(),
416                        reason: e.to_string(),
417                    }
418                })?;
419                Ok(Some(offset_info))
420            }
421            None => Ok(None),
422        }
423    }
424
425    async fn get_all_offsets(
426        &self,
427        consumer_group: &str,
428    ) -> StateResult<HashMap<TopicPartition, OffsetInfo>> {
429        let prefix = self.build_group_prefix(consumer_group);
430        let keys = self.backend.list_keys(&prefix).await?;
431
432        let mut offsets = HashMap::new();
433
434        for key in keys {
435            if let Some(value) = self.backend.get(&key).await? {
436                if let Ok(offset_info) = bincode::deserialize::<OffsetInfo>(&value) {
437                    // Extract topic-partition from key
438                    let key_str = String::from_utf8_lossy(&key);
439                    if let Some(tp_str) = key_str.strip_prefix(&format!("{}{}:", self.prefix, consumer_group)) {
440                        if let Some((topic, partition_str)) = tp_str.rsplit_once(':') {
441                            if let Ok(partition) = partition_str.parse::<u32>() {
442                                let tp = TopicPartition::new(topic.to_string(), partition);
443                                offsets.insert(tp, offset_info);
444                            }
445                        }
446                    }
447                }
448            }
449        }
450
451        Ok(offsets)
452    }
453
454    async fn delete_offset(
455        &self,
456        consumer_group: &str,
457        topic_partition: &TopicPartition,
458    ) -> StateResult<()> {
459        let key = self.build_key(consumer_group, topic_partition);
460        self.backend.delete(&key).await
461    }
462
463    async fn clear_offsets(&self, consumer_group: &str) -> StateResult<()> {
464        let prefix = self.build_group_prefix(consumer_group);
465        let keys = self.backend.list_keys(&prefix).await?;
466
467        for key in keys {
468            self.backend.delete(&key).await?;
469        }
470
471        Ok(())
472    }
473}
474
475/// Main offset manager
476///
477/// Coordinates offset tracking, committing, and recovery for Kafka consumers.
478///
479/// # Features
480///
481/// - Per-partition offset tracking
482/// - Watermark-based commit strategies
483/// - Checkpoint integration
484/// - Offset lag monitoring
485/// - Automatic and manual commit modes
486pub struct OffsetManager<B: StateBackend, S: OffsetStore> {
487    /// Consumer group ID
488    consumer_group: String,
489    /// State backend for checkpointing
490    backend: Arc<B>,
491    /// Offset storage
492    store: Arc<S>,
493    /// Commit strategy
494    strategy: OffsetCommitStrategy,
495    /// Pending offsets (not yet committed)
496    pending_offsets: Arc<DashMap<TopicPartition, OffsetInfo>>,
497    /// Committed offsets
498    committed_offsets: Arc<DashMap<TopicPartition, OffsetInfo>>,
499    /// Statistics
500    stats: Arc<RwLock<OffsetStats>>,
501    /// Last commit time
502    last_commit: Arc<Mutex<DateTime<Utc>>>,
503    /// Event counter for EveryNEvents strategy
504    event_counter: Arc<Mutex<usize>>,
505    /// Watermark tracker for OnWatermark strategy
506    last_watermark: Arc<Mutex<Option<Watermark>>>,
507}
508
509impl<B: StateBackend, S: OffsetStore> OffsetManager<B, S> {
510    /// Create a new offset manager
511    ///
512    /// # Arguments
513    /// * `consumer_group` - Consumer group ID
514    /// * `backend` - State backend for checkpointing
515    /// * `store` - Offset storage implementation
516    /// * `strategy` - Commit strategy
517    pub fn new(
518        consumer_group: String,
519        backend: Arc<B>,
520        store: Arc<S>,
521        strategy: OffsetCommitStrategy,
522    ) -> Self {
523        Self {
524            consumer_group,
525            backend,
526            store,
527            strategy,
528            pending_offsets: Arc::new(DashMap::new()),
529            committed_offsets: Arc::new(DashMap::new()),
530            stats: Arc::new(RwLock::new(OffsetStats::default())),
531            last_commit: Arc::new(Mutex::new(Utc::now())),
532            event_counter: Arc::new(Mutex::new(0)),
533            last_watermark: Arc::new(Mutex::new(None)),
534        }
535    }
536
537    /// Update offset for a partition
538    ///
539    /// # Arguments
540    /// * `topic` - Kafka topic
541    /// * `partition` - Partition number
542    /// * `offset` - New offset value
543    ///
544    /// # Returns
545    /// * `Ok(())` - Offset updated
546    /// * `Err(StateError)` - Update failed
547    pub async fn update_offset(
548        &mut self,
549        topic: &str,
550        partition: u32,
551        offset: i64,
552    ) -> StateResult<()> {
553        let tp = TopicPartition::new(topic.to_string(), partition);
554        let offset_info = OffsetInfo::new(offset);
555
556        self.pending_offsets.insert(tp.clone(), offset_info);
557
558        debug!(
559            consumer_group = %self.consumer_group,
560            topic = topic,
561            partition = partition,
562            offset = offset,
563            "Updated offset"
564        );
565
566        // Check if we should auto-commit
567        if self.should_auto_commit().await {
568            self.commit().await?;
569        }
570
571        Ok(())
572    }
573
574    /// Update offset with high watermark
575    pub async fn update_offset_with_watermark(
576        &mut self,
577        topic: &str,
578        partition: u32,
579        offset: i64,
580        high_watermark: i64,
581    ) -> StateResult<()> {
582        let tp = TopicPartition::new(topic.to_string(), partition);
583        let offset_info = OffsetInfo::with_high_watermark(offset, high_watermark);
584
585        self.pending_offsets.insert(tp.clone(), offset_info);
586
587        debug!(
588            consumer_group = %self.consumer_group,
589            topic = topic,
590            partition = partition,
591            offset = offset,
592            high_watermark = high_watermark,
593            "Updated offset with high watermark"
594        );
595
596        // Check if we should auto-commit
597        if self.should_auto_commit().await {
598            self.commit().await?;
599        }
600
601        Ok(())
602    }
603
604    /// Get current offset for a partition
605    ///
606    /// Returns the committed offset, or None if no offset has been committed.
607    pub async fn get_offset(&self, topic: &str, partition: u32) -> StateResult<Option<i64>> {
608        let tp = TopicPartition::new(topic.to_string(), partition);
609
610        // Check committed offsets first
611        if let Some(info) = self.committed_offsets.get(&tp) {
612            return Ok(Some(info.offset));
613        }
614
615        // Fall back to stored offsets
616        if let Some(info) = self.store.get_offset(&self.consumer_group, &tp).await? {
617            return Ok(Some(info.offset));
618        }
619
620        Ok(None)
621    }
622
623    /// Get offset info (includes metadata and watermarks)
624    pub async fn get_offset_info(
625        &self,
626        topic: &str,
627        partition: u32,
628    ) -> StateResult<Option<OffsetInfo>> {
629        let tp = TopicPartition::new(topic.to_string(), partition);
630
631        // Check committed offsets first
632        if let Some(info) = self.committed_offsets.get(&tp) {
633            return Ok(Some(info.clone()));
634        }
635
636        // Fall back to stored offsets
637        self.store.get_offset(&self.consumer_group, &tp).await
638    }
639
640    /// Commit all pending offsets
641    ///
642    /// # Returns
643    /// * `Ok(committed_count)` - Number of offsets committed
644    /// * `Err(StateError)` - Commit failed
645    pub async fn commit(&mut self) -> StateResult<usize> {
646        let start_time = Utc::now();
647        let pending: Vec<_> = self
648            .pending_offsets
649            .iter()
650            .map(|entry| (entry.key().clone(), entry.value().clone()))
651            .collect();
652
653        if pending.is_empty() {
654            debug!(consumer_group = %self.consumer_group, "No pending offsets to commit");
655            return Ok(0);
656        }
657
658        info!(
659            consumer_group = %self.consumer_group,
660            count = pending.len(),
661            "Committing offsets"
662        );
663
664        let mut stats = self.stats.write().await;
665        stats.total_commits += 1;
666
667        let mut committed_count = 0;
668        let mut failed = false;
669
670        for (tp, offset_info) in pending {
671            match self
672                .store
673                .store_offset(&self.consumer_group, &tp, &offset_info)
674                .await
675            {
676                Ok(()) => {
677                    self.committed_offsets.insert(tp.clone(), offset_info.clone());
678                    self.pending_offsets.remove(&tp);
679                    committed_count += 1;
680                    debug!(topic_partition = %tp, offset = offset_info.offset, "Committed offset");
681                }
682                Err(e) => {
683                    error!(
684                        topic_partition = %tp,
685                        error = %e,
686                        "Failed to commit offset"
687                    );
688                    failed = true;
689                }
690            }
691        }
692
693        // Update statistics
694        if failed {
695            stats.failed_commits += 1;
696        } else {
697            stats.successful_commits += 1;
698        }
699        stats.total_offsets_committed += committed_count as u64;
700        stats.last_commit_time = Some(start_time);
701
702        let duration = Utc::now().signed_duration_since(start_time);
703        let duration_ms = duration.num_milliseconds() as u64;
704
705        // Update rolling average
706        if stats.avg_commit_duration_ms == 0 {
707            stats.avg_commit_duration_ms = duration_ms;
708        } else {
709            stats.avg_commit_duration_ms =
710                (stats.avg_commit_duration_ms * 9 + duration_ms) / 10;
711        }
712
713        drop(stats);
714
715        *self.last_commit.lock().await = Utc::now();
716
717        info!(
718            consumer_group = %self.consumer_group,
719            committed = committed_count,
720            duration_ms = duration_ms,
721            "Offset commit completed"
722        );
723
724        if failed {
725            Err(StateError::TransactionFailed {
726                operation: "offset_commit".to_string(),
727                reason: format!("Failed to commit some offsets, {} succeeded", committed_count),
728            })
729        } else {
730            Ok(committed_count)
731        }
732    }
733
734    /// Calculate offset lag for a partition
735    ///
736    /// # Arguments
737    /// * `topic` - Kafka topic
738    /// * `partition` - Partition number
739    /// * `high_watermark` - Current high watermark (latest offset in partition)
740    ///
741    /// # Returns
742    /// * `Ok(lag)` - Offset lag (high_watermark - current_offset)
743    /// * `Err(StateError)` - Calculation failed
744    pub async fn lag(
745        &self,
746        topic: &str,
747        partition: u32,
748        high_watermark: i64,
749    ) -> StateResult<i64> {
750        let current_offset = self.get_offset(topic, partition).await?.unwrap_or(0);
751        Ok(high_watermark - current_offset)
752    }
753
754    /// Get total lag across all partitions
755    pub async fn total_lag(&self) -> StateResult<i64> {
756        let mut total = 0i64;
757
758        for entry in self.committed_offsets.iter() {
759            if let Some(lag) = entry.value().lag() {
760                total += lag;
761            }
762        }
763
764        Ok(total)
765    }
766
767    /// Restore offsets from storage
768    ///
769    /// Loads all committed offsets for the consumer group from the offset store.
770    ///
771    /// # Returns
772    /// * `Ok(restored_count)` - Number of offsets restored
773    /// * `Err(StateError)` - Restore failed
774    pub async fn restore(&mut self) -> StateResult<usize> {
775        info!(consumer_group = %self.consumer_group, "Restoring offsets");
776
777        let offsets = self.store.get_all_offsets(&self.consumer_group).await?;
778        let count = offsets.len();
779
780        for (tp, offset_info) in offsets {
781            self.committed_offsets.insert(tp.clone(), offset_info);
782            debug!(topic_partition = %tp, "Restored offset");
783        }
784
785        info!(
786            consumer_group = %self.consumer_group,
787            count = count,
788            "Offset restore completed"
789        );
790
791        Ok(count)
792    }
793
794    /// Reset offset for a partition
795    ///
796    /// # Arguments
797    /// * `topic` - Kafka topic
798    /// * `partition` - Partition number
799    /// * `strategy` - Reset strategy (earliest, latest, specific offset, etc.)
800    /// * `partition_high_watermark` - Current high watermark for the partition
801    ///
802    /// # Returns
803    /// * `Ok(new_offset)` - The new offset value
804    /// * `Err(StateError)` - Reset failed
805    pub async fn reset_offset(
806        &mut self,
807        topic: &str,
808        partition: u32,
809        strategy: OffsetResetStrategy,
810        partition_high_watermark: Option<i64>,
811    ) -> StateResult<i64> {
812        let tp = TopicPartition::new(topic.to_string(), partition);
813
814        let new_offset = match strategy {
815            OffsetResetStrategy::Earliest => 0,
816            OffsetResetStrategy::Latest => partition_high_watermark.unwrap_or(0),
817            OffsetResetStrategy::Specific(offset) => offset,
818            OffsetResetStrategy::Timestamp(ts) => {
819                // In a real implementation, this would query Kafka for offset by timestamp
820                // For now, we'll just use the timestamp value as offset
821                warn!(
822                    topic = topic,
823                    partition = partition,
824                    timestamp = ts,
825                    "Timestamp-based reset not fully implemented, using timestamp as offset"
826                );
827                ts
828            }
829        };
830
831        let offset_info = OffsetInfo::new(new_offset);
832        self.store
833            .store_offset(&self.consumer_group, &tp, &offset_info)
834            .await?;
835        self.committed_offsets.insert(tp.clone(), offset_info);
836        self.pending_offsets.remove(&tp);
837
838        info!(
839            consumer_group = %self.consumer_group,
840            topic = topic,
841            partition = partition,
842            new_offset = new_offset,
843            strategy = ?strategy,
844            "Reset offset"
845        );
846
847        Ok(new_offset)
848    }
849
850    /// Check if automatic commit should occur based on strategy
851    async fn should_auto_commit(&mut self) -> bool {
852        match &self.strategy {
853            OffsetCommitStrategy::Periodic { interval } => {
854                let last_commit = *self.last_commit.lock().await;
855                let elapsed = Utc::now().signed_duration_since(last_commit);
856                elapsed.num_milliseconds() as u64 >= interval.as_millis() as u64
857            }
858            OffsetCommitStrategy::EveryNEvents { n } => {
859                let mut counter = self.event_counter.lock().await;
860                *counter += 1;
861                if *counter >= *n {
862                    *counter = 0;
863                    true
864                } else {
865                    false
866                }
867            }
868            OffsetCommitStrategy::OnWatermark => false, // Handled externally via notify_watermark
869            OffsetCommitStrategy::Manual => false,
870        }
871    }
872
873    /// Notify manager of watermark advancement (for OnWatermark strategy)
874    ///
875    /// # Arguments
876    /// * `watermark` - New watermark value
877    ///
878    /// # Returns
879    /// * `Ok(committed)` - True if commit occurred
880    /// * `Err(StateError)` - Notification failed
881    pub async fn notify_watermark(&mut self, watermark: Watermark) -> StateResult<bool> {
882        let mut last_wm = self.last_watermark.lock().await;
883
884        if matches!(self.strategy, OffsetCommitStrategy::OnWatermark) {
885            if last_wm.is_none() || watermark > last_wm.unwrap() {
886                *last_wm = Some(watermark);
887                drop(last_wm);
888
889                self.commit().await?;
890                return Ok(true);
891            }
892        }
893
894        Ok(false)
895    }
896
897    /// Get statistics
898    pub async fn stats(&self) -> OffsetStats {
899        self.stats.read().await.clone()
900    }
901
902    /// Get all committed offsets
903    pub async fn get_all_committed(&self) -> HashMap<TopicPartition, OffsetInfo> {
904        self.committed_offsets
905            .iter()
906            .map(|entry| (entry.key().clone(), entry.value().clone()))
907            .collect()
908    }
909
910    /// Get all pending offsets (not yet committed)
911    pub async fn get_all_pending(&self) -> HashMap<TopicPartition, OffsetInfo> {
912        self.pending_offsets
913            .iter()
914            .map(|entry| (entry.key().clone(), entry.value().clone()))
915            .collect()
916    }
917
918    /// Clear all offsets for this consumer group
919    pub async fn clear(&mut self) -> StateResult<()> {
920        info!(consumer_group = %self.consumer_group, "Clearing all offsets");
921
922        self.store.clear_offsets(&self.consumer_group).await?;
923        self.committed_offsets.clear();
924        self.pending_offsets.clear();
925
926        Ok(())
927    }
928}
929
930#[cfg(test)]
931mod tests {
932    use super::*;
933    use crate::state::MemoryStateBackend;
934
935    #[test]
936    fn test_topic_partition() {
937        let tp = TopicPartition::new("test-topic".to_string(), 5);
938        assert_eq!(tp.topic, "test-topic");
939        assert_eq!(tp.partition, 5);
940        assert_eq!(tp.to_key(), "test-topic:5");
941        assert_eq!(tp.to_string(), "test-topic:5");
942    }
943
944    #[test]
945    fn test_offset_info_lag() {
946        let info = OffsetInfo::with_high_watermark(100, 200);
947        assert_eq!(info.lag(), Some(100));
948
949        let info_no_hwm = OffsetInfo::new(100);
950        assert_eq!(info_no_hwm.lag(), None);
951    }
952
953    #[tokio::test]
954    async fn test_in_memory_offset_store() {
955        let store = InMemoryOffsetStore::new();
956        let tp = TopicPartition::new("test".to_string(), 0);
957        let offset_info = OffsetInfo::new(100);
958
959        // Store offset
960        store
961            .store_offset("group1", &tp, &offset_info)
962            .await
963            .unwrap();
964
965        // Retrieve offset
966        let retrieved = store.get_offset("group1", &tp).await.unwrap();
967        assert!(retrieved.is_some());
968        assert_eq!(retrieved.unwrap().offset, 100);
969
970        // Get all offsets
971        let all = store.get_all_offsets("group1").await.unwrap();
972        assert_eq!(all.len(), 1);
973
974        // Delete offset
975        store.delete_offset("group1", &tp).await.unwrap();
976        let retrieved = store.get_offset("group1", &tp).await.unwrap();
977        assert!(retrieved.is_none());
978    }
979
980    #[tokio::test]
981    async fn test_state_backend_offset_store() {
982        let backend = Arc::new(MemoryStateBackend::new(None));
983        let store = StateBackendOffsetStore::new(backend.clone());
984        let tp = TopicPartition::new("test-topic".to_string(), 3);
985        let offset_info = OffsetInfo::new(500);
986
987        // Store offset
988        store
989            .store_offset("group1", &tp, &offset_info)
990            .await
991            .unwrap();
992
993        // Retrieve offset
994        let retrieved = store.get_offset("group1", &tp).await.unwrap();
995        assert!(retrieved.is_some());
996        assert_eq!(retrieved.unwrap().offset, 500);
997
998        // Get all offsets
999        let all = store.get_all_offsets("group1").await.unwrap();
1000        assert_eq!(all.len(), 1);
1001        assert_eq!(all.get(&tp).unwrap().offset, 500);
1002
1003        // Clear offsets
1004        store.clear_offsets("group1").await.unwrap();
1005        let all = store.get_all_offsets("group1").await.unwrap();
1006        assert_eq!(all.len(), 0);
1007    }
1008
1009    #[tokio::test]
1010    async fn test_offset_manager_update_and_commit() {
1011        let backend = Arc::new(MemoryStateBackend::new(None));
1012        let store = Arc::new(InMemoryOffsetStore::new());
1013        let mut manager = OffsetManager::new(
1014            "test-group".to_string(),
1015            backend,
1016            store.clone(),
1017            OffsetCommitStrategy::Manual,
1018        );
1019
1020        // Update offset
1021        manager.update_offset("topic1", 0, 100).await.unwrap();
1022
1023        // Offset should be pending
1024        let pending = manager.get_all_pending().await;
1025        assert_eq!(pending.len(), 1);
1026
1027        // Commit
1028        let committed = manager.commit().await.unwrap();
1029        assert_eq!(committed, 1);
1030
1031        // Should be in committed offsets now
1032        let offset = manager.get_offset("topic1", 0).await.unwrap();
1033        assert_eq!(offset, Some(100));
1034
1035        // Pending should be empty
1036        let pending = manager.get_all_pending().await;
1037        assert_eq!(pending.len(), 0);
1038    }
1039
1040    #[tokio::test]
1041    async fn test_offset_manager_lag() {
1042        let backend = Arc::new(MemoryStateBackend::new(None));
1043        let store = Arc::new(InMemoryOffsetStore::new());
1044        let mut manager = OffsetManager::new(
1045            "test-group".to_string(),
1046            backend,
1047            store,
1048            OffsetCommitStrategy::Manual,
1049        );
1050
1051        manager.update_offset("topic1", 0, 100).await.unwrap();
1052        manager.commit().await.unwrap();
1053
1054        let lag = manager.lag("topic1", 0, 200).await.unwrap();
1055        assert_eq!(lag, 100);
1056    }
1057
1058    #[tokio::test]
1059    async fn test_offset_manager_restore() {
1060        let backend = Arc::new(MemoryStateBackend::new(None));
1061        let store = Arc::new(InMemoryOffsetStore::new());
1062
1063        // Create first manager and commit some offsets
1064        let mut manager1 = OffsetManager::new(
1065            "test-group".to_string(),
1066            backend.clone(),
1067            store.clone(),
1068            OffsetCommitStrategy::Manual,
1069        );
1070
1071        manager1.update_offset("topic1", 0, 100).await.unwrap();
1072        manager1.update_offset("topic1", 1, 200).await.unwrap();
1073        manager1.commit().await.unwrap();
1074
1075        // Create second manager and restore
1076        let mut manager2 = OffsetManager::new(
1077            "test-group".to_string(),
1078            backend,
1079            store,
1080            OffsetCommitStrategy::Manual,
1081        );
1082
1083        let restored = manager2.restore().await.unwrap();
1084        assert_eq!(restored, 2);
1085
1086        assert_eq!(manager2.get_offset("topic1", 0).await.unwrap(), Some(100));
1087        assert_eq!(manager2.get_offset("topic1", 1).await.unwrap(), Some(200));
1088    }
1089
1090    #[tokio::test]
1091    async fn test_offset_manager_reset() {
1092        let backend = Arc::new(MemoryStateBackend::new(None));
1093        let store = Arc::new(InMemoryOffsetStore::new());
1094        let mut manager = OffsetManager::new(
1095            "test-group".to_string(),
1096            backend,
1097            store,
1098            OffsetCommitStrategy::Manual,
1099        );
1100
1101        manager.update_offset("topic1", 0, 100).await.unwrap();
1102        manager.commit().await.unwrap();
1103
1104        // Reset to earliest
1105        let new_offset = manager
1106            .reset_offset("topic1", 0, OffsetResetStrategy::Earliest, None)
1107            .await
1108            .unwrap();
1109        assert_eq!(new_offset, 0);
1110        assert_eq!(manager.get_offset("topic1", 0).await.unwrap(), Some(0));
1111
1112        // Reset to latest
1113        let new_offset = manager
1114            .reset_offset("topic1", 0, OffsetResetStrategy::Latest, Some(500))
1115            .await
1116            .unwrap();
1117        assert_eq!(new_offset, 500);
1118        assert_eq!(manager.get_offset("topic1", 0).await.unwrap(), Some(500));
1119
1120        // Reset to specific
1121        let new_offset = manager
1122            .reset_offset("topic1", 0, OffsetResetStrategy::Specific(250), None)
1123            .await
1124            .unwrap();
1125        assert_eq!(new_offset, 250);
1126        assert_eq!(manager.get_offset("topic1", 0).await.unwrap(), Some(250));
1127    }
1128
1129    #[tokio::test]
1130    async fn test_offset_manager_periodic_strategy() {
1131        let backend = Arc::new(MemoryStateBackend::new(None));
1132        let store = Arc::new(InMemoryOffsetStore::new());
1133        let mut manager = OffsetManager::new(
1134            "test-group".to_string(),
1135            backend,
1136            store,
1137            OffsetCommitStrategy::Periodic {
1138                interval: Duration::from_millis(100),
1139            },
1140        );
1141
1142        // First update - should not auto-commit yet
1143        manager.update_offset("topic1", 0, 100).await.unwrap();
1144        assert_eq!(manager.get_all_pending().await.len(), 1);
1145
1146        // Wait for interval
1147        tokio::time::sleep(Duration::from_millis(150)).await;
1148
1149        // Second update - should auto-commit
1150        manager.update_offset("topic1", 0, 200).await.unwrap();
1151
1152        // Previous offset should be committed
1153        assert_eq!(manager.get_offset("topic1", 0).await.unwrap(), Some(200));
1154    }
1155
1156    #[tokio::test]
1157    async fn test_offset_manager_watermark_strategy() {
1158        let backend = Arc::new(MemoryStateBackend::new(None));
1159        let store = Arc::new(InMemoryOffsetStore::new());
1160        let mut manager = OffsetManager::new(
1161            "test-group".to_string(),
1162            backend,
1163            store,
1164            OffsetCommitStrategy::OnWatermark,
1165        );
1166
1167        manager.update_offset("topic1", 0, 100).await.unwrap();
1168        assert_eq!(manager.get_all_pending().await.len(), 1);
1169
1170        // Notify watermark advancement
1171        let watermark = Watermark::new(1000);
1172        let committed = manager.notify_watermark(watermark).await.unwrap();
1173        assert!(committed);
1174
1175        // Should be committed now
1176        assert_eq!(manager.get_offset("topic1", 0).await.unwrap(), Some(100));
1177        assert_eq!(manager.get_all_pending().await.len(), 0);
1178    }
1179
1180    #[tokio::test]
1181    async fn test_offset_manager_stats() {
1182        let backend = Arc::new(MemoryStateBackend::new(None));
1183        let store = Arc::new(InMemoryOffsetStore::new());
1184        let mut manager = OffsetManager::new(
1185            "test-group".to_string(),
1186            backend,
1187            store,
1188            OffsetCommitStrategy::Manual,
1189        );
1190
1191        let stats = manager.stats().await;
1192        assert_eq!(stats.total_commits, 0);
1193
1194        manager.update_offset("topic1", 0, 100).await.unwrap();
1195        manager.commit().await.unwrap();
1196
1197        let stats = manager.stats().await;
1198        assert_eq!(stats.total_commits, 1);
1199        assert_eq!(stats.successful_commits, 1);
1200        assert_eq!(stats.total_offsets_committed, 1);
1201        assert!(stats.last_commit_time.is_some());
1202    }
1203
1204    #[tokio::test]
1205    async fn test_offset_manager_multiple_partitions() {
1206        let backend = Arc::new(MemoryStateBackend::new(None));
1207        let store = Arc::new(InMemoryOffsetStore::new());
1208        let mut manager = OffsetManager::new(
1209            "test-group".to_string(),
1210            backend,
1211            store,
1212            OffsetCommitStrategy::Manual,
1213        );
1214
1215        // Update multiple partitions
1216        manager.update_offset("topic1", 0, 100).await.unwrap();
1217        manager.update_offset("topic1", 1, 200).await.unwrap();
1218        manager.update_offset("topic2", 0, 300).await.unwrap();
1219
1220        let pending = manager.get_all_pending().await;
1221        assert_eq!(pending.len(), 3);
1222
1223        // Commit all
1224        let committed = manager.commit().await.unwrap();
1225        assert_eq!(committed, 3);
1226
1227        // Verify all committed
1228        assert_eq!(manager.get_offset("topic1", 0).await.unwrap(), Some(100));
1229        assert_eq!(manager.get_offset("topic1", 1).await.unwrap(), Some(200));
1230        assert_eq!(manager.get_offset("topic2", 0).await.unwrap(), Some(300));
1231    }
1232
1233    #[tokio::test]
1234    async fn test_offset_manager_with_high_watermark() {
1235        let backend = Arc::new(MemoryStateBackend::new(None));
1236        let store = Arc::new(InMemoryOffsetStore::new());
1237        let mut manager = OffsetManager::new(
1238            "test-group".to_string(),
1239            backend,
1240            store,
1241            OffsetCommitStrategy::Manual,
1242        );
1243
1244        manager
1245            .update_offset_with_watermark("topic1", 0, 100, 200)
1246            .await
1247            .unwrap();
1248        manager.commit().await.unwrap();
1249
1250        let info = manager.get_offset_info("topic1", 0).await.unwrap();
1251        assert!(info.is_some());
1252        let info = info.unwrap();
1253        assert_eq!(info.offset, 100);
1254        assert_eq!(info.high_watermark, Some(200));
1255        assert_eq!(info.lag(), Some(100));
1256
1257        let total_lag = manager.total_lag().await.unwrap();
1258        assert_eq!(total_lag, 100);
1259    }
1260}
llm_optimizer_processor/kafka/offset_manager.rs

llm_optimizer_processor/kafka/
offset_manager.rs
