soma_application/provider_registry/
generations.rs

1use super::*;
2
3const GENERATION_HISTORY_LIMIT: usize = 3;
4
5#[derive(Clone)]
6pub(super) struct StoredGeneration {
7    pub(super) generation_id: u64,
8    pub(super) providers: BTreeMap<String, Arc<dyn Provider>>,
9    pub(super) core_registry: CoreRegistry,
10    pub(super) snapshot: Arc<RegistrySnapshot>,
11    pub(super) file_fingerprint: Option<String>,
12    pub(super) python_environments: PythonProviderEnvironmentSelections,
13}
14
15pub(super) struct GenerationTransition {
16    suspended: Vec<Arc<dyn Provider>>,
17    evicted: Vec<Arc<dyn Provider>>,
18}
19
20impl ProviderRegistry {
21    /// Returns the active generation and bounded rollback history.
22    pub fn python_generation_status(&self) -> Value {
23        let state = self
24            .state
25            .read()
26            .expect("provider registry lock should not be poisoned");
27        let describe = |generation_id: u64,
28                        snapshot: &Arc<RegistrySnapshot>,
29                        file_fingerprint: &Option<String>,
30                        environments: &PythonProviderEnvironmentSelections| {
31            json!({
32                "generation_id": generation_id,
33                "snapshot_id": snapshot.id,
34                "catalog_fingerprint": snapshot.fingerprint,
35                "file_fingerprint": file_fingerprint,
36                "python_environment_count": environments.len(),
37            })
38        };
39        json!({
40            "active": describe(
41                state.generation_id,
42                &state.snapshot,
43                &state.file_fingerprint,
44                &state.python_environments,
45            ),
46            "rollback_candidates": state.history.iter().map(|generation| {
47                describe(
48                    generation.generation_id,
49                    &generation.snapshot,
50                    &generation.file_fingerprint,
51                    &generation.python_environments,
52                )
53            }).collect::<Vec<_>>(),
54        })
55    }
56
57    /// Atomically reactivates a retained immutable provider generation.
58    pub async fn rollback_python_generation(
59        &self,
60        target_generation_id: u64,
61    ) -> Result<Value, ProviderError> {
62        let target_environments = {
63            let state = self
64                .state
65                .read()
66                .expect("provider registry lock should not be poisoned");
67            state
68                .history
69                .iter()
70                .find(|generation| generation.generation_id == target_generation_id)
71                .map(|generation| generation.python_environments.clone())
72                .ok_or_else(|| unavailable_generation(target_generation_id))?
73        };
74        let rollback_fingerprint = if let Some(file_source) = self.file_source.clone() {
75            let environments = target_environments.clone();
76            Some(
77                tokio::task::spawn_blocking(move || {
78                    file_source.fingerprint_with_python_environments(&environments)
79                })
80                .await
81                .map_err(fingerprint_error)?
82                .map_err(fingerprint_error)?,
83            )
84        } else {
85            None
86        };
87        let (report, transition) = {
88            let mut state = self
89                .state
90                .write()
91                .expect("provider registry lock should not be poisoned");
92            let position = state
93                .history
94                .iter()
95                .position(|generation| generation.generation_id == target_generation_id)
96                .ok_or_else(|| unavailable_generation(target_generation_id))?;
97            let target = state
98                .history
99                .remove(position)
100                .expect("located generation must remain present");
101            let previous_generation_id = state.generation_id;
102            let transition = publish_generation(
103                &mut state,
104                target.providers,
105                target.core_registry,
106                target.snapshot,
107                rollback_fingerprint,
108                target.python_environments,
109            );
110            (
111                json!({
112                    "previous_generation_id": previous_generation_id,
113                    "restored_generation_id": target_generation_id,
114                    "active_generation_id": state.generation_id,
115                    "snapshot_id": state.snapshot.id,
116                }),
117                transition,
118            )
119        };
120        settle_transition_async(transition).await;
121        Ok(report)
122    }
123}
124
125fn unavailable_generation(target_generation_id: u64) -> ProviderError {
126    ProviderError::validation(
127        "registry",
128        "python_generation_rollback",
129        "python_generation_unavailable",
130        format!("generation `{target_generation_id}` is not in the rollback window"),
131    )
132}
133
134fn fingerprint_error(error: impl std::fmt::Display) -> ProviderError {
135    ProviderError::new(
136        "provider_file_fingerprint_failed",
137        "registry",
138        Some("python_generation_rollback".to_owned()),
139        error.to_string(),
140        "Inspect the provider directory and retry.",
141    )
142}
143
144pub(super) fn publish_generation(
145    state: &mut RegistryState,
146    providers: BTreeMap<String, Arc<dyn Provider>>,
147    core_registry: CoreRegistry,
148    snapshot: Arc<RegistrySnapshot>,
149    file_fingerprint: Option<String>,
150    python_environments: PythonProviderEnvironmentSelections,
151) -> GenerationTransition {
152    let suspended: Vec<_> = state
153        .providers
154        .values()
155        .filter(|provider| {
156            !providers
157                .values()
158                .any(|candidate| Arc::ptr_eq(candidate, provider))
159        })
160        .cloned()
161        .collect();
162    for provider in &suspended {
163        provider.deactivate();
164    }
165    for provider in providers.values() {
166        provider.activate();
167    }
168    state.history.push_front(StoredGeneration {
169        generation_id: state.generation_id,
170        providers: state.providers.clone(),
171        core_registry: state.core_registry.clone(),
172        snapshot: state.snapshot.clone(),
173        file_fingerprint: state.file_fingerprint.clone(),
174        python_environments: state.python_environments.clone(),
175    });
176    state.generation_id = state.generation_id.saturating_add(1);
177    state.providers = providers;
178    state.core_registry = core_registry;
179    state.snapshot = snapshot;
180    state.file_fingerprint = file_fingerprint;
181    state.python_environments = python_environments;
182
183    let mut evicted = Vec::new();
184    while state.history.len() > GENERATION_HISTORY_LIMIT {
185        if let Some(generation) = state.history.pop_back() {
186            for provider in generation.providers.into_values() {
187                let retained = state
188                    .providers
189                    .values()
190                    .chain(
191                        state
192                            .history
193                            .iter()
194                            .flat_map(|generation| generation.providers.values()),
195                    )
196                    .any(|candidate| Arc::ptr_eq(candidate, &provider));
197                if !retained {
198                    evicted.push(provider);
199                }
200            }
201        }
202    }
203    GenerationTransition { suspended, evicted }
204}
205
206pub(super) fn settle_transition(transition: GenerationTransition) {
207    if let Ok(runtime) = tokio::runtime::Handle::try_current() {
208        runtime.spawn(settle_transition_async(transition));
209        return;
210    }
211    tokio::runtime::Builder::new_current_thread()
212        .enable_all()
213        .build()
214        .expect("provider transition runtime should build")
215        .block_on(settle_transition_async(transition));
216}
217
218pub(super) async fn settle_transition_async(transition: GenerationTransition) {
219    for provider in transition.suspended {
220        provider.suspend().await;
221    }
222    for provider in transition.evicted {
223        provider.retire().await;
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use std::sync::atomic::{AtomicUsize, Ordering};
230
231    use super::*;
232
233    struct LifecycleProbe {
234        suspended: AtomicUsize,
235        retired: AtomicUsize,
236    }
237
238    #[async_trait]
239    impl Provider for LifecycleProbe {
240        fn catalog(&self) -> ProviderCatalog {
241            panic!("lifecycle settlement does not inspect catalogs")
242        }
243
244        async fn call(&self, _call: ProviderCall) -> Result<ProviderOutput, ProviderError> {
245            panic!("lifecycle settlement does not dispatch")
246        }
247
248        async fn suspend(&self) {
249            self.suspended.fetch_add(1, Ordering::SeqCst);
250        }
251
252        async fn retire(&self) {
253            self.retired.fetch_add(1, Ordering::SeqCst);
254        }
255    }
256
257    #[test]
258    fn synchronous_transition_settles_without_an_ambient_tokio_runtime() {
259        let probe = Arc::new(LifecycleProbe {
260            suspended: AtomicUsize::new(0),
261            retired: AtomicUsize::new(0),
262        });
263        settle_transition(GenerationTransition {
264            suspended: vec![probe.clone()],
265            evicted: vec![probe.clone()],
266        });
267        assert_eq!(probe.suspended.load(Ordering::SeqCst), 1);
268        assert_eq!(probe.retired.load(Ordering::SeqCst), 1);
269    }
270}
soma_application/provider_registry/generations.rs

soma_application/provider_registry/
generations.rs
