prodigy/cook/execution/mapreduce/checkpoint/
environment.rs

1//! Checkpoint environment for Reader pattern effects
2//!
3//! This module provides the CheckpointEnv type and Reader pattern helpers
4//! for accessing checkpoint-related components in Effect-based code.
5
6use super::effects::storage::CheckpointStorageEnv;
7use super::pure::triggers::CheckpointTriggerConfig;
8use super::{CheckpointStorage, MapReduceCheckpoint};
9use chrono::{DateTime, Utc};
10use std::path::PathBuf;
11use std::sync::atomic::{AtomicUsize, Ordering};
12use std::sync::Arc;
13use stillwater::{asks, local, Effect};
14use tokio::sync::RwLock;
15
16/// Environment for checkpoint operations
17///
18/// This environment provides all dependencies needed for checkpoint operations
19/// via the Reader pattern. Effects can access these through helper functions.
20#[derive(Clone)]
21pub struct CheckpointEnv {
22    /// Job identifier
23    pub job_id: String,
24    /// Checkpoint storage implementation
25    pub storage: Arc<dyn CheckpointStorage>,
26    /// Current checkpoint state
27    pub current_checkpoint: Arc<RwLock<Option<MapReduceCheckpoint>>>,
28    /// Storage path for checkpoints
29    pub storage_path: PathBuf,
30    /// Checkpoint trigger configuration
31    pub trigger_config: CheckpointTriggerConfig,
32    /// Items processed since last checkpoint
33    pub items_since_checkpoint: Arc<AtomicUsize>,
34    /// Time of last checkpoint
35    pub last_checkpoint_time: Arc<RwLock<DateTime<Utc>>>,
36    /// Whether checkpointing is enabled
37    pub enabled: bool,
38}
39
40impl CheckpointEnv {
41    /// Create a new checkpoint environment
42    pub fn new(
43        job_id: String,
44        storage: Arc<dyn CheckpointStorage>,
45        storage_path: PathBuf,
46        trigger_config: CheckpointTriggerConfig,
47    ) -> Self {
48        Self {
49            job_id,
50            storage,
51            current_checkpoint: Arc::new(RwLock::new(None)),
52            storage_path,
53            trigger_config,
54            items_since_checkpoint: Arc::new(AtomicUsize::new(0)),
55            last_checkpoint_time: Arc::new(RwLock::new(Utc::now())),
56            enabled: true,
57        }
58    }
59
60    /// Create a disabled checkpoint environment (for testing)
61    pub fn disabled() -> Self {
62        use super::FileCheckpointStorage;
63
64        let temp_path = std::env::temp_dir().join("prodigy_disabled_checkpoints");
65        let storage: Arc<dyn CheckpointStorage> =
66            Arc::new(FileCheckpointStorage::new(temp_path.clone(), false));
67
68        Self {
69            job_id: "disabled".to_string(),
70            storage,
71            current_checkpoint: Arc::new(RwLock::new(None)),
72            storage_path: temp_path,
73            trigger_config: CheckpointTriggerConfig::none(),
74            items_since_checkpoint: Arc::new(AtomicUsize::new(0)),
75            last_checkpoint_time: Arc::new(RwLock::new(Utc::now())),
76            enabled: false,
77        }
78    }
79
80    /// Increment items processed count
81    pub fn increment_items(&self, count: usize) {
82        self.items_since_checkpoint
83            .fetch_add(count, Ordering::SeqCst);
84    }
85
86    /// Reset items processed count
87    pub fn reset_items(&self) {
88        self.items_since_checkpoint.store(0, Ordering::SeqCst);
89    }
90
91    /// Get current items processed count
92    pub fn get_items(&self) -> usize {
93        self.items_since_checkpoint.load(Ordering::Acquire)
94    }
95}
96
97// Implement CheckpointStorageEnv for CheckpointEnv
98impl CheckpointStorageEnv for CheckpointEnv {
99    fn storage(&self) -> Arc<dyn CheckpointStorage> {
100        Arc::clone(&self.storage)
101    }
102
103    fn current_checkpoint(&self) -> Arc<RwLock<Option<MapReduceCheckpoint>>> {
104        Arc::clone(&self.current_checkpoint)
105    }
106
107    fn storage_path(&self) -> PathBuf {
108        self.storage_path.clone()
109    }
110}
111
112// =============================================================================
113// Reader Pattern Helpers
114// =============================================================================
115
116/// Error type for checkpoint operations
117#[derive(Debug, Clone, thiserror::Error)]
118pub enum CheckpointError {
119    #[error("Checkpointing is disabled")]
120    Disabled,
121
122    #[error("No checkpoint available")]
123    NoCheckpoint,
124
125    #[error("Storage error: {0}")]
126    Storage(String),
127
128    #[error("Validation error: {0}")]
129    Validation(String),
130}
131
132/// Get the job ID from the checkpoint environment.
133pub fn get_checkpoint_job_id(
134) -> impl Effect<Output = String, Error = CheckpointError, Env = CheckpointEnv> {
135    asks(|env: &CheckpointEnv| env.job_id.clone())
136}
137
138/// Get the checkpoint trigger configuration.
139pub fn get_trigger_config(
140) -> impl Effect<Output = CheckpointTriggerConfig, Error = CheckpointError, Env = CheckpointEnv> {
141    asks(|env: &CheckpointEnv| env.trigger_config.clone())
142}
143
144/// Get the checkpoint storage.
145pub fn get_checkpoint_storage(
146) -> impl Effect<Output = Arc<dyn CheckpointStorage>, Error = CheckpointError, Env = CheckpointEnv>
147{
148    asks(|env: &CheckpointEnv| env.storage.clone())
149}
150
151/// Get items processed since last checkpoint.
152pub fn get_items_since_checkpoint(
153) -> impl Effect<Output = usize, Error = CheckpointError, Env = CheckpointEnv> {
154    asks(|env: &CheckpointEnv| env.items_since_checkpoint.load(Ordering::Acquire))
155}
156
157/// Check if checkpointing is enabled.
158pub fn is_checkpointing_enabled(
159) -> impl Effect<Output = bool, Error = CheckpointError, Env = CheckpointEnv> {
160    asks(|env: &CheckpointEnv| env.enabled)
161}
162
163/// Get the storage path.
164pub fn get_checkpoint_storage_path(
165) -> impl Effect<Output = PathBuf, Error = CheckpointError, Env = CheckpointEnv> {
166    asks(|env: &CheckpointEnv| env.storage_path.clone())
167}
168
169// =============================================================================
170// Local Override Utilities
171// =============================================================================
172
173/// Run an effect with checkpointing disabled.
174pub fn with_checkpointing_disabled<E>(
175    effect: E,
176) -> impl Effect<Output = E::Output, Error = CheckpointError, Env = CheckpointEnv>
177where
178    E: Effect<Error = CheckpointError, Env = CheckpointEnv>,
179{
180    local(
181        |env: &CheckpointEnv| CheckpointEnv {
182            enabled: false,
183            ..env.clone()
184        },
185        effect,
186    )
187}
188
189/// Run an effect with a custom trigger configuration.
190pub fn with_trigger_config<E>(
191    config: CheckpointTriggerConfig,
192    effect: E,
193) -> impl Effect<Output = E::Output, Error = CheckpointError, Env = CheckpointEnv>
194where
195    E: Effect<Error = CheckpointError, Env = CheckpointEnv>,
196{
197    local(
198        move |env: &CheckpointEnv| CheckpointEnv {
199            trigger_config: config.clone(),
200            ..env.clone()
201        },
202        effect,
203    )
204}
205
206// =============================================================================
207// Mock Environment Builder
208// =============================================================================
209
210/// Builder for creating mock CheckpointEnv instances for testing.
211#[derive(Clone)]
212pub struct MockCheckpointEnvBuilder {
213    job_id: String,
214    trigger_config: CheckpointTriggerConfig,
215    enabled: bool,
216    initial_checkpoint: Option<MapReduceCheckpoint>,
217}
218
219impl Default for MockCheckpointEnvBuilder {
220    fn default() -> Self {
221        Self::new()
222    }
223}
224
225impl MockCheckpointEnvBuilder {
226    /// Create a new builder with default values.
227    pub fn new() -> Self {
228        Self {
229            job_id: "mock-job-123".to_string(),
230            trigger_config: CheckpointTriggerConfig::default(),
231            enabled: true,
232            initial_checkpoint: None,
233        }
234    }
235
236    /// Set the job ID.
237    pub fn with_job_id(mut self, job_id: impl Into<String>) -> Self {
238        self.job_id = job_id.into();
239        self
240    }
241
242    /// Set the trigger configuration.
243    pub fn with_trigger_config(mut self, config: CheckpointTriggerConfig) -> Self {
244        self.trigger_config = config;
245        self
246    }
247
248    /// Disable checkpointing.
249    pub fn disabled(mut self) -> Self {
250        self.enabled = false;
251        self
252    }
253
254    /// Set initial checkpoint.
255    pub fn with_checkpoint(mut self, checkpoint: MapReduceCheckpoint) -> Self {
256        self.initial_checkpoint = Some(checkpoint);
257        self
258    }
259
260    /// Build the mock environment.
261    pub fn build(self) -> CheckpointEnv {
262        use super::FileCheckpointStorage;
263
264        let temp_dir = std::env::temp_dir().join(format!("prodigy_mock_{}", self.job_id));
265        let _ = std::fs::create_dir_all(&temp_dir);
266        let storage: Arc<dyn CheckpointStorage> =
267            Arc::new(FileCheckpointStorage::new(temp_dir.clone(), true));
268
269        let mut env = CheckpointEnv::new(self.job_id, storage, temp_dir, self.trigger_config);
270        env.enabled = self.enabled;
271
272        if let Some(checkpoint) = self.initial_checkpoint {
273            let current_checkpoint = Arc::clone(&env.current_checkpoint);
274            tokio::task::block_in_place(|| {
275                tokio::runtime::Handle::current().block_on(async {
276                    *current_checkpoint.write().await = Some(checkpoint);
277                })
278            });
279        }
280
281        env
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    #[tokio::test]
290    async fn test_get_checkpoint_job_id() {
291        let env = MockCheckpointEnvBuilder::new()
292            .with_job_id("my-test-job")
293            .build();
294
295        let effect = get_checkpoint_job_id();
296        let result = effect.run(&env).await;
297
298        assert!(result.is_ok());
299        assert_eq!(result.unwrap(), "my-test-job");
300    }
301
302    #[tokio::test]
303    async fn test_get_trigger_config() {
304        let config = CheckpointTriggerConfig::item_interval(10);
305        let env = MockCheckpointEnvBuilder::new()
306            .with_trigger_config(config)
307            .build();
308
309        let effect = get_trigger_config();
310        let result = effect.run(&env).await;
311
312        assert!(result.is_ok());
313        assert_eq!(result.unwrap().agent_completion_interval, Some(10));
314    }
315
316    #[tokio::test]
317    async fn test_is_checkpointing_enabled() {
318        let enabled_env = MockCheckpointEnvBuilder::new().build();
319        let disabled_env = MockCheckpointEnvBuilder::new().disabled().build();
320
321        // Test enabled environment
322        let effect = is_checkpointing_enabled();
323        assert!(effect.run(&enabled_env).await.unwrap());
324
325        // Test disabled environment separately
326        let effect = is_checkpointing_enabled();
327        assert!(!effect.run(&disabled_env).await.unwrap());
328    }
329
330    #[tokio::test]
331    async fn test_with_checkpointing_disabled() {
332        let env = MockCheckpointEnvBuilder::new().build();
333
334        // Normally enabled
335        let effect = is_checkpointing_enabled();
336        assert!(effect.run(&env).await.unwrap());
337
338        // Disabled via local override
339        let effect = with_checkpointing_disabled(is_checkpointing_enabled());
340        assert!(!effect.run(&env).await.unwrap());
341
342        // Original unchanged
343        let effect = is_checkpointing_enabled();
344        assert!(effect.run(&env).await.unwrap());
345    }
346
347    #[tokio::test]
348    async fn test_with_trigger_config_override() {
349        let env = MockCheckpointEnvBuilder::new()
350            .with_trigger_config(CheckpointTriggerConfig::item_interval(5))
351            .build();
352
353        // Without override
354        let effect = get_trigger_config();
355        assert_eq!(
356            effect.run(&env).await.unwrap().agent_completion_interval,
357            Some(5)
358        );
359
360        // With override
361        let new_config = CheckpointTriggerConfig::item_interval(100);
362        let effect = with_trigger_config(new_config, get_trigger_config());
363        assert_eq!(
364            effect.run(&env).await.unwrap().agent_completion_interval,
365            Some(100)
366        );
367
368        // Original unchanged
369        let effect = get_trigger_config();
370        assert_eq!(
371            effect.run(&env).await.unwrap().agent_completion_interval,
372            Some(5)
373        );
374    }
375
376    #[test]
377    fn test_checkpoint_env_increment_items() {
378        let env = MockCheckpointEnvBuilder::new().build();
379
380        assert_eq!(env.get_items(), 0);
381        env.increment_items(5);
382        assert_eq!(env.get_items(), 5);
383        env.increment_items(3);
384        assert_eq!(env.get_items(), 8);
385        env.reset_items();
386        assert_eq!(env.get_items(), 0);
387    }
388}
prodigy/cook/execution/mapreduce/checkpoint/environment.rs

prodigy/cook/execution/mapreduce/checkpoint/
environment.rs
