llm_optimizer_processor/state/
checkpoint.rs

1//! Checkpointing and recovery for state backends
2//!
3//! This module provides checkpoint coordination for creating periodic snapshots
4//! of state backends. Checkpoints enable fault tolerance and allow state to be
5//! recovered after failures.
6
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11use std::time::Duration;
12use tokio::sync::{Mutex, RwLock};
13use tokio::task::JoinHandle;
14use tracing::{debug, error, info, warn};
15
16use super::backend::StateBackend;
17use crate::error::{StateError, StateResult};
18
19/// Checkpoint metadata
20///
21/// Contains information about a checkpoint including when it was created,
22/// how many entries it contains, and whether it's been validated.
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct CheckpointMetadata {
25    /// Unique checkpoint identifier
26    pub checkpoint_id: String,
27    /// When the checkpoint was created
28    pub created_at: DateTime<Utc>,
29    /// Number of key-value pairs in the checkpoint
30    pub entry_count: usize,
31    /// Size of checkpoint file in bytes
32    pub size_bytes: u64,
33    /// Checksum for integrity verification (SHA-256)
34    pub checksum: String,
35    /// Whether this checkpoint has been validated
36    pub validated: bool,
37    /// Version of the checkpoint format
38    pub version: u32,
39}
40
41impl CheckpointMetadata {
42    /// Create new checkpoint metadata
43    pub fn new(checkpoint_id: String, entry_count: usize, size_bytes: u64, checksum: String) -> Self {
44        Self {
45            checkpoint_id,
46            created_at: Utc::now(),
47            entry_count,
48            size_bytes,
49            checksum,
50            validated: false,
51            version: 1,
52        }
53    }
54}
55
56/// A checkpoint containing state data and metadata
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct Checkpoint {
59    /// Checkpoint metadata
60    pub metadata: CheckpointMetadata,
61    /// The actual state data (key-value pairs)
62    pub data: Vec<(Vec<u8>, Vec<u8>)>,
63}
64
65impl Checkpoint {
66    /// Create a new checkpoint
67    pub fn new(checkpoint_id: String, data: Vec<(Vec<u8>, Vec<u8>)>) -> Self {
68        let size_bytes = data
69            .iter()
70            .map(|(k, v)| k.len() + v.len())
71            .sum::<usize>() as u64;
72
73        let checksum = Self::calculate_checksum(&data);
74
75        Self {
76            metadata: CheckpointMetadata::new(
77                checkpoint_id,
78                data.len(),
79                size_bytes,
80                checksum,
81            ),
82            data,
83        }
84    }
85
86    /// Calculate SHA-256 checksum of the data
87    fn calculate_checksum(data: &[(Vec<u8>, Vec<u8>)]) -> String {
88        use std::collections::hash_map::DefaultHasher;
89        use std::hash::{Hash, Hasher};
90
91        let mut hasher = DefaultHasher::new();
92        for (key, value) in data {
93            key.hash(&mut hasher);
94            value.hash(&mut hasher);
95        }
96
97        format!("{:x}", hasher.finish())
98    }
99
100    /// Validate the checkpoint's integrity
101    pub fn validate(&mut self) -> StateResult<()> {
102        let calculated_checksum = Self::calculate_checksum(&self.data);
103
104        if calculated_checksum != self.metadata.checksum {
105            return Err(StateError::CheckpointFailed {
106                checkpoint_id: self.metadata.checkpoint_id.clone(),
107                reason: format!(
108                    "Checksum mismatch: expected {}, got {}",
109                    self.metadata.checksum, calculated_checksum
110                ),
111            });
112        }
113
114        self.metadata.validated = true;
115        Ok(())
116    }
117
118    /// Save checkpoint to a file
119    pub async fn save<P: AsRef<Path>>(&self, path: P) -> StateResult<()> {
120        let path = path.as_ref();
121        info!("Saving checkpoint {} to {:?}", self.metadata.checkpoint_id, path);
122
123        // Create parent directory if needed
124        if let Some(parent) = path.parent() {
125            tokio::fs::create_dir_all(parent).await.map_err(|e| {
126                StateError::CheckpointFailed {
127                    checkpoint_id: self.metadata.checkpoint_id.clone(),
128                    reason: format!("Failed to create directory: {}", e),
129                }
130            })?;
131        }
132
133        // Serialize checkpoint
134        let serialized = bincode::serialize(self).map_err(|e| StateError::SerializationFailed {
135            key: self.metadata.checkpoint_id.clone(),
136            reason: e.to_string(),
137        })?;
138
139        // Write to file
140        tokio::fs::write(path, &serialized).await.map_err(|e| {
141            StateError::CheckpointFailed {
142                checkpoint_id: self.metadata.checkpoint_id.clone(),
143                reason: format!("Failed to write file: {}", e),
144            }
145        })?;
146
147        info!(
148            "Checkpoint saved: {} entries, {} bytes",
149            self.metadata.entry_count, self.metadata.size_bytes
150        );
151
152        Ok(())
153    }
154
155    /// Load checkpoint from a file
156    pub async fn load<P: AsRef<Path>>(path: P) -> StateResult<Self> {
157        let path = path.as_ref();
158        info!("Loading checkpoint from {:?}", path);
159
160        // Read file
161        let data = tokio::fs::read(path).await.map_err(|e| {
162            StateError::RestoreFailed {
163                checkpoint_id: path.to_string_lossy().to_string(),
164                reason: format!("Failed to read file: {}", e),
165            }
166        })?;
167
168        // Deserialize
169        let mut checkpoint: Checkpoint =
170            bincode::deserialize(&data).map_err(|e| StateError::DeserializationFailed {
171                key: path.to_string_lossy().to_string(),
172                reason: e.to_string(),
173            })?;
174
175        // Validate
176        checkpoint.validate()?;
177
178        info!(
179            "Checkpoint loaded: {} entries, {} bytes",
180            checkpoint.metadata.entry_count, checkpoint.metadata.size_bytes
181        );
182
183        Ok(checkpoint)
184    }
185}
186
187/// Options for checkpoint creation
188#[derive(Debug, Clone)]
189pub struct CheckpointOptions {
190    /// Checkpoint directory
191    pub checkpoint_dir: PathBuf,
192    /// Maximum number of checkpoints to keep
193    pub max_checkpoints: usize,
194    /// Whether to compress checkpoints
195    pub compress: bool,
196    /// Minimum interval between checkpoints
197    pub min_interval: Duration,
198}
199
200impl Default for CheckpointOptions {
201    fn default() -> Self {
202        Self {
203            checkpoint_dir: PathBuf::from("/tmp/checkpoints"),
204            max_checkpoints: 5,
205            compress: true,
206            min_interval: Duration::from_secs(60),
207        }
208    }
209}
210
211/// Statistics about checkpointing
212#[derive(Debug, Clone, Default)]
213pub struct CheckpointStats {
214    /// Total number of checkpoints created
215    pub checkpoints_created: u64,
216    /// Total number of checkpoint failures
217    pub checkpoint_failures: u64,
218    /// Total number of restores
219    pub restores: u64,
220    /// Total number of restore failures
221    pub restore_failures: u64,
222    /// Last checkpoint timestamp
223    pub last_checkpoint_time: Option<DateTime<Utc>>,
224    /// Last checkpoint duration in milliseconds
225    pub last_checkpoint_duration_ms: Option<u64>,
226    /// Total bytes checkpointed
227    pub total_bytes_checkpointed: u64,
228}
229
230/// Checkpoint coordinator
231///
232/// Manages periodic checkpoint creation and restoration for a state backend.
233///
234/// ## Features
235///
236/// - **Periodic checkpoints**: Automatically creates checkpoints at regular intervals
237/// - **Retention policy**: Keeps only the N most recent checkpoints
238/// - **Graceful shutdown**: Ensures final checkpoint on shutdown
239/// - **Recovery**: Can restore from the latest valid checkpoint
240///
241/// ## Example
242///
243/// ```rust,no_run
244/// use processor::state::{
245///     CheckpointCoordinator, CheckpointOptions, MemoryStateBackend
246/// };
247/// use std::sync::Arc;
248/// use std::time::Duration;
249/// use std::path::PathBuf;
250///
251/// #[tokio::main]
252/// async fn main() -> anyhow::Result<()> {
253///     let backend = Arc::new(MemoryStateBackend::new(None));
254///
255///     let options = CheckpointOptions {
256///         checkpoint_dir: PathBuf::from("/var/lib/optimizer/checkpoints"),
257///         max_checkpoints: 10,
258///         compress: true,
259///         min_interval: Duration::from_secs(300),
260///     };
261///
262///     let coordinator = CheckpointCoordinator::new(
263///         backend.clone(),
264///         Duration::from_secs(300),
265///         options,
266///     );
267///
268///     // Start periodic checkpointing
269///     coordinator.start().await?;
270///
271///     // Work with state...
272///
273///     // Shutdown with final checkpoint
274///     coordinator.shutdown().await?;
275///
276///     Ok(())
277/// }
278/// ```
279pub struct CheckpointCoordinator<B: StateBackend> {
280    /// The state backend to checkpoint
281    backend: Arc<B>,
282    /// Checkpoint interval
283    interval: Duration,
284    /// Checkpoint options
285    options: CheckpointOptions,
286    /// Statistics
287    stats: Arc<RwLock<CheckpointStats>>,
288    /// Background task handle
289    task_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
290    /// Shutdown signal
291    shutdown_tx: Arc<Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
292}
293
294impl<B: StateBackend + 'static> CheckpointCoordinator<B> {
295    /// Create a new checkpoint coordinator
296    ///
297    /// # Arguments
298    ///
299    /// * `backend` - The state backend to checkpoint
300    /// * `interval` - How often to create checkpoints
301    /// * `options` - Checkpoint configuration options
302    ///
303    /// # Returns
304    ///
305    /// A new coordinator instance
306    pub fn new(backend: Arc<B>, interval: Duration, options: CheckpointOptions) -> Self {
307        Self {
308            backend,
309            interval,
310            options,
311            stats: Arc::new(RwLock::new(CheckpointStats::default())),
312            task_handle: Arc::new(Mutex::new(None)),
313            shutdown_tx: Arc::new(Mutex::new(None)),
314        }
315    }
316
317    /// Start periodic checkpointing
318    ///
319    /// Spawns a background task that creates checkpoints at the configured interval.
320    ///
321    /// # Returns
322    ///
323    /// * `Ok(())` - Checkpointing started
324    /// * `Err(StateError)` - Failed to start
325    pub async fn start(&self) -> StateResult<()> {
326        let mut handle_guard = self.task_handle.lock().await;
327
328        if handle_guard.is_some() {
329            warn!("Checkpoint coordinator already started");
330            return Ok(());
331        }
332
333        info!("Starting checkpoint coordinator with interval {:?}", self.interval);
334
335        let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
336        *self.shutdown_tx.lock().await = Some(shutdown_tx);
337
338        let backend = Arc::clone(&self.backend);
339        let options = self.options.clone();
340        let stats = Arc::clone(&self.stats);
341        let interval = self.interval;
342
343        let handle = tokio::spawn(async move {
344            let mut interval_timer = tokio::time::interval(interval);
345            let mut last_checkpoint = Utc::now();
346
347            loop {
348                tokio::select! {
349                    _ = interval_timer.tick() => {
350                        // Check minimum interval
351                        let elapsed = Utc::now().signed_duration_since(last_checkpoint);
352                        if elapsed.num_milliseconds() as u64 >= options.min_interval.as_millis() as u64 {
353                            if let Err(e) = Self::create_checkpoint_internal(
354                                &backend,
355                                &options,
356                                &stats,
357                            ).await {
358                                error!("Checkpoint failed: {}", e);
359                                let mut stats_guard = stats.write().await;
360                                stats_guard.checkpoint_failures += 1;
361                            } else {
362                                last_checkpoint = Utc::now();
363                            }
364                        }
365                    }
366                    _ = &mut shutdown_rx => {
367                        info!("Checkpoint coordinator shutting down");
368                        // Create final checkpoint
369                        if let Err(e) = Self::create_checkpoint_internal(
370                            &backend,
371                            &options,
372                            &stats,
373                        ).await {
374                            error!("Final checkpoint failed: {}", e);
375                        }
376                        break;
377                    }
378                }
379            }
380        });
381
382        *handle_guard = Some(handle);
383        Ok(())
384    }
385
386    /// Create a checkpoint immediately
387    ///
388    /// # Returns
389    ///
390    /// * `Ok(checkpoint_id)` - Checkpoint created successfully
391    /// * `Err(StateError)` - Checkpoint failed
392    pub async fn create_checkpoint(&self) -> StateResult<String> {
393        Self::create_checkpoint_internal(&self.backend, &self.options, &self.stats).await
394    }
395
396    /// Internal checkpoint creation
397    async fn create_checkpoint_internal(
398        backend: &Arc<B>,
399        options: &CheckpointOptions,
400        stats: &Arc<RwLock<CheckpointStats>>,
401    ) -> StateResult<String> {
402        let start_time = Utc::now();
403        let checkpoint_id = format!("checkpoint_{}", start_time.timestamp_millis());
404
405        debug!("Creating checkpoint {}", checkpoint_id);
406
407        // Get all state data
408        let data = Self::export_state(backend).await?;
409
410        // Create checkpoint
411        let checkpoint = Checkpoint::new(checkpoint_id.clone(), data);
412
413        // Save to disk
414        let checkpoint_path = options
415            .checkpoint_dir
416            .join(format!("{}.ckpt", checkpoint_id));
417
418        checkpoint.save(&checkpoint_path).await?;
419
420        // Update statistics
421        let duration = Utc::now().signed_duration_since(start_time);
422        let mut stats_guard = stats.write().await;
423        stats_guard.checkpoints_created += 1;
424        stats_guard.last_checkpoint_time = Some(start_time);
425        stats_guard.last_checkpoint_duration_ms = Some(duration.num_milliseconds() as u64);
426        stats_guard.total_bytes_checkpointed += checkpoint.metadata.size_bytes;
427        drop(stats_guard);
428
429        // Clean up old checkpoints
430        Self::cleanup_old_checkpoints(options).await?;
431
432        info!(
433            "Checkpoint {} created in {}ms",
434            checkpoint_id,
435            duration.num_milliseconds()
436        );
437
438        Ok(checkpoint_id)
439    }
440
441    /// Export all state from the backend
442    async fn export_state(backend: &Arc<B>) -> StateResult<Vec<(Vec<u8>, Vec<u8>)>> {
443        let keys = backend.list_keys(b"").await?;
444        let mut data = Vec::with_capacity(keys.len());
445
446        for key in keys {
447            if let Some(value) = backend.get(&key).await? {
448                data.push((key, value));
449            }
450        }
451
452        Ok(data)
453    }
454
455    /// Clean up old checkpoints based on retention policy
456    async fn cleanup_old_checkpoints(options: &CheckpointOptions) -> StateResult<()> {
457        let mut entries = tokio::fs::read_dir(&options.checkpoint_dir)
458            .await
459            .map_err(|e| StateError::StorageError {
460                backend_type: "checkpoint".to_string(),
461                details: format!("Failed to read checkpoint directory: {}", e),
462            })?;
463
464        let mut checkpoints = Vec::new();
465
466        while let Some(entry) = entries.next_entry().await.map_err(|e| {
467            StateError::StorageError {
468                backend_type: "checkpoint".to_string(),
469                details: format!("Failed to read directory entry: {}", e),
470            }
471        })? {
472            let path = entry.path();
473            if path.extension().and_then(|s| s.to_str()) == Some("ckpt") {
474                if let Ok(metadata) = entry.metadata().await {
475                    if let Ok(modified) = metadata.modified() {
476                        checkpoints.push((path, modified));
477                    }
478                }
479            }
480        }
481
482        // Sort by modification time (newest first)
483        checkpoints.sort_by(|a, b| b.1.cmp(&a.1));
484
485        // Delete old checkpoints
486        for (path, _) in checkpoints.into_iter().skip(options.max_checkpoints) {
487            debug!("Deleting old checkpoint: {:?}", path);
488            if let Err(e) = tokio::fs::remove_file(&path).await {
489                warn!("Failed to delete checkpoint {:?}: {}", path, e);
490            }
491        }
492
493        Ok(())
494    }
495
496    /// Restore from the latest checkpoint
497    ///
498    /// # Returns
499    ///
500    /// * `Ok(checkpoint_id)` - Successfully restored
501    /// * `Err(StateError)` - No checkpoint found or restore failed
502    pub async fn restore_latest(&self) -> StateResult<String> {
503        info!("Restoring from latest checkpoint");
504
505        // Find latest checkpoint
506        let checkpoint_path = self.find_latest_checkpoint().await?;
507
508        // Load checkpoint
509        let checkpoint = Checkpoint::load(&checkpoint_path).await?;
510
511        // Restore to backend
512        self.restore_checkpoint(&checkpoint).await?;
513
514        let mut stats = self.stats.write().await;
515        stats.restores += 1;
516
517        Ok(checkpoint.metadata.checkpoint_id)
518    }
519
520    /// Find the latest checkpoint file
521    async fn find_latest_checkpoint(&self) -> StateResult<PathBuf> {
522        let mut entries = tokio::fs::read_dir(&self.options.checkpoint_dir)
523            .await
524            .map_err(|e| StateError::RestoreFailed {
525                checkpoint_id: "unknown".to_string(),
526                reason: format!("Failed to read checkpoint directory: {}", e),
527            })?;
528
529        let mut latest: Option<(PathBuf, std::time::SystemTime)> = None;
530
531        while let Some(entry) = entries.next_entry().await.map_err(|e| {
532            StateError::RestoreFailed {
533                checkpoint_id: "unknown".to_string(),
534                reason: format!("Failed to read directory entry: {}", e),
535            }
536        })? {
537            let path = entry.path();
538            if path.extension().and_then(|s| s.to_str()) == Some("ckpt") {
539                if let Ok(metadata) = entry.metadata().await {
540                    if let Ok(modified) = metadata.modified() {
541                        if latest.is_none() || modified > latest.as_ref().unwrap().1 {
542                            latest = Some((path, modified));
543                        }
544                    }
545                }
546            }
547        }
548
549        latest.map(|(path, _)| path).ok_or_else(|| StateError::RestoreFailed {
550            checkpoint_id: "unknown".to_string(),
551            reason: "No checkpoint files found".to_string(),
552        })
553    }
554
555    /// Restore state from a checkpoint
556    async fn restore_checkpoint(&self, checkpoint: &Checkpoint) -> StateResult<()> {
557        info!(
558            "Restoring checkpoint {} with {} entries",
559            checkpoint.metadata.checkpoint_id, checkpoint.metadata.entry_count
560        );
561
562        // Clear existing state
563        self.backend.clear().await?;
564
565        // Restore all entries
566        for (key, value) in &checkpoint.data {
567            self.backend.put(key, value).await?;
568        }
569
570        info!("Checkpoint restored successfully");
571        Ok(())
572    }
573
574    /// Get checkpoint statistics
575    ///
576    /// # Returns
577    ///
578    /// Current checkpoint statistics
579    pub async fn stats(&self) -> CheckpointStats {
580        self.stats.read().await.clone()
581    }
582
583    /// Shutdown the coordinator
584    ///
585    /// Stops the background task and creates a final checkpoint.
586    ///
587    /// # Returns
588    ///
589    /// * `Ok(())` - Shutdown successful
590    /// * `Err(StateError)` - Shutdown failed
591    pub async fn shutdown(&self) -> StateResult<()> {
592        info!("Shutting down checkpoint coordinator");
593
594        // Send shutdown signal
595        if let Some(tx) = self.shutdown_tx.lock().await.take() {
596            let _ = tx.send(());
597        }
598
599        // Wait for background task
600        if let Some(handle) = self.task_handle.lock().await.take() {
601            if let Err(e) = handle.await {
602                error!("Background task panicked: {}", e);
603            }
604        }
605
606        info!("Checkpoint coordinator shutdown complete");
607        Ok(())
608    }
609}
610
611#[cfg(test)]
612mod tests {
613    use super::*;
614    use crate::state::MemoryStateBackend;
615
616    #[tokio::test]
617    async fn test_checkpoint_creation() {
618        let backend = Arc::new(MemoryStateBackend::new(None));
619
620        backend.put(b"key1", b"value1").await.unwrap();
621        backend.put(b"key2", b"value2").await.unwrap();
622
623        let data = vec![
624            (b"key1".to_vec(), b"value1".to_vec()),
625            (b"key2".to_vec(), b"value2".to_vec()),
626        ];
627
628        let checkpoint = Checkpoint::new("test_checkpoint".to_string(), data);
629        assert_eq!(checkpoint.metadata.entry_count, 2);
630        assert!(checkpoint.metadata.size_bytes > 0);
631    }
632
633    #[tokio::test]
634    async fn test_checkpoint_validation() {
635        let data = vec![(b"key1".to_vec(), b"value1".to_vec())];
636        let mut checkpoint = Checkpoint::new("test".to_string(), data);
637
638        checkpoint.validate().unwrap();
639        assert!(checkpoint.metadata.validated);
640
641        // Corrupt the data
642        checkpoint.data.push((b"key2".to_vec(), b"value2".to_vec()));
643
644        assert!(checkpoint.validate().is_err());
645    }
646
647    #[tokio::test]
648    async fn test_checkpoint_save_load() {
649        let temp_dir = std::env::temp_dir().join(format!("checkpoint_test_{}", uuid::Uuid::new_v4()));
650        tokio::fs::create_dir_all(&temp_dir).await.unwrap();
651
652        let data = vec![
653            (b"key1".to_vec(), b"value1".to_vec()),
654            (b"key2".to_vec(), b"value2".to_vec()),
655        ];
656
657        let checkpoint = Checkpoint::new("test".to_string(), data);
658        let checkpoint_path = temp_dir.join("test.ckpt");
659
660        checkpoint.save(&checkpoint_path).await.unwrap();
661        assert!(checkpoint_path.exists());
662
663        let loaded = Checkpoint::load(&checkpoint_path).await.unwrap();
664        assert_eq!(loaded.metadata.entry_count, 2);
665        assert!(loaded.metadata.validated);
666
667        // Cleanup
668        tokio::fs::remove_dir_all(&temp_dir).await.ok();
669    }
670
671    #[tokio::test]
672    async fn test_coordinator_create_checkpoint() {
673        let temp_dir = std::env::temp_dir().join(format!("coordinator_test_{}", uuid::Uuid::new_v4()));
674
675        let backend = Arc::new(MemoryStateBackend::new(None));
676        backend.put(b"key1", b"value1").await.unwrap();
677        backend.put(b"key2", b"value2").await.unwrap();
678
679        let options = CheckpointOptions {
680            checkpoint_dir: temp_dir.clone(),
681            max_checkpoints: 3,
682            compress: false,
683            min_interval: Duration::from_secs(1),
684        };
685
686        let coordinator = CheckpointCoordinator::new(
687            backend.clone(),
688            Duration::from_secs(60),
689            options,
690        );
691
692        let checkpoint_id = coordinator.create_checkpoint().await.unwrap();
693        assert!(!checkpoint_id.is_empty());
694
695        let stats = coordinator.stats().await;
696        assert_eq!(stats.checkpoints_created, 1);
697
698        // Cleanup
699        tokio::fs::remove_dir_all(&temp_dir).await.ok();
700    }
701
702    #[tokio::test]
703    async fn test_coordinator_restore() {
704        let temp_dir = std::env::temp_dir().join(format!("restore_test_{}", uuid::Uuid::new_v4()));
705
706        let backend = Arc::new(MemoryStateBackend::new(None));
707        backend.put(b"key1", b"value1").await.unwrap();
708        backend.put(b"key2", b"value2").await.unwrap();
709
710        let options = CheckpointOptions {
711            checkpoint_dir: temp_dir.clone(),
712            max_checkpoints: 3,
713            compress: false,
714            min_interval: Duration::from_secs(1),
715        };
716
717        let coordinator = CheckpointCoordinator::new(
718            backend.clone(),
719            Duration::from_secs(60),
720            options,
721        );
722
723        // Create checkpoint
724        coordinator.create_checkpoint().await.unwrap();
725
726        // Clear backend
727        backend.clear().await.unwrap();
728        assert_eq!(backend.count().await.unwrap(), 0);
729
730        // Restore
731        coordinator.restore_latest().await.unwrap();
732        assert_eq!(backend.count().await.unwrap(), 2);
733        assert_eq!(
734            backend.get(b"key1").await.unwrap(),
735            Some(b"value1".to_vec())
736        );
737
738        // Cleanup
739        tokio::fs::remove_dir_all(&temp_dir).await.ok();
740    }
741
742    #[tokio::test]
743    async fn test_checkpoint_retention() {
744        let temp_dir = std::env::temp_dir().join(format!("retention_test_{}", uuid::Uuid::new_v4()));
745
746        let backend = Arc::new(MemoryStateBackend::new(None));
747
748        let options = CheckpointOptions {
749            checkpoint_dir: temp_dir.clone(),
750            max_checkpoints: 3,
751            compress: false,
752            min_interval: Duration::from_millis(100),
753        };
754
755        let coordinator = CheckpointCoordinator::new(
756            backend.clone(),
757            Duration::from_secs(60),
758            options,
759        );
760
761        // Create multiple checkpoints
762        for i in 0..5 {
763            backend
764                .put(&format!("key{}", i).into_bytes(), b"value")
765                .await
766                .unwrap();
767            coordinator.create_checkpoint().await.unwrap();
768            tokio::time::sleep(Duration::from_millis(200)).await;
769        }
770
771        // Count checkpoint files
772        let mut count = 0;
773        let mut entries = tokio::fs::read_dir(&temp_dir).await.unwrap();
774        while let Some(_entry) = entries.next_entry().await.unwrap() {
775            count += 1;
776        }
777
778        // Should only keep max_checkpoints
779        assert!(count <= 3);
780
781        // Cleanup
782        tokio::fs::remove_dir_all(&temp_dir).await.ok();
783    }
784}
llm_optimizer_processor/state/checkpoint.rs

llm_optimizer_processor/state/
checkpoint.rs
