awa_model/
storage.rs

1use crate::admin;
2use crate::error::AwaError;
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use sqlx::{FromRow, PgExecutor, PgPool};
6use std::collections::BTreeMap;
7
8#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)]
9pub struct StorageStatus {
10    pub current_engine: String,
11    pub active_engine: String,
12    pub prepared_engine: Option<String>,
13    pub state: String,
14    pub transition_epoch: i64,
15    pub details: serde_json::Value,
16    pub entered_at: DateTime<Utc>,
17    pub updated_at: DateTime<Utc>,
18    pub finalized_at: Option<DateTime<Utc>>,
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
22pub struct StorageStatusReport {
23    #[serde(flatten)]
24    pub status: StorageStatus,
25    pub canonical_live_backlog: i64,
26    pub prepared_queue_storage_schema: Option<String>,
27    pub prepared_schema_ready: bool,
28    pub live_runtime_capability_counts: BTreeMap<String, usize>,
29    pub can_enter_mixed_transition: bool,
30    pub enter_mixed_transition_blockers: Vec<String>,
31    pub can_finalize: bool,
32    pub finalize_blockers: Vec<String>,
33}
34
35pub async fn status<'e, E>(executor: E) -> Result<StorageStatus, AwaError>
36where
37    E: PgExecutor<'e>,
38{
39    sqlx::query_as::<_, StorageStatus>("SELECT * FROM awa.storage_status()")
40        .fetch_one(executor)
41        .await
42        .map_err(AwaError::from)
43}
44
45pub async fn prepare<'e, E>(
46    executor: E,
47    engine: &str,
48    details: serde_json::Value,
49) -> Result<StorageStatus, AwaError>
50where
51    E: PgExecutor<'e>,
52{
53    sqlx::query_as::<_, StorageStatus>("SELECT * FROM awa.storage_prepare($1, $2)")
54        .bind(engine)
55        .bind(details)
56        .fetch_one(executor)
57        .await
58        .map_err(AwaError::from)
59}
60
61pub async fn abort<'e, E>(executor: E) -> Result<StorageStatus, AwaError>
62where
63    E: PgExecutor<'e>,
64{
65    sqlx::query_as::<_, StorageStatus>("SELECT * FROM awa.storage_abort()")
66        .fetch_one(executor)
67        .await
68        .map_err(AwaError::from)
69}
70
71fn queue_storage_schema_from_status(status: &StorageStatus) -> Option<String> {
72    let queue_storage_involved = status.prepared_engine.as_deref() == Some("queue_storage")
73        || status.active_engine == "queue_storage";
74    if !queue_storage_involved {
75        return None;
76    }
77    Some(
78        status
79            .details
80            .get("schema")
81            .and_then(serde_json::Value::as_str)
82            .unwrap_or("awa_exp")
83            .to_string(),
84    )
85}
86
87async fn prepared_queue_storage_schema_ready(
88    pool: &PgPool,
89    schema: &str,
90) -> Result<bool, AwaError> {
91    sqlx::query_scalar::<_, bool>(
92        r#"
93        SELECT
94            to_regclass(format('%I.%I', $1, 'queue_ring_state')) IS NOT NULL
95            AND to_regclass(format('%I.%I', $1, 'ready_entries')) IS NOT NULL
96            AND to_regclass(format('%I.%I', $1, 'leases')) IS NOT NULL
97        "#,
98    )
99    .bind(schema)
100    .fetch_one(pool)
101    .await
102    .map_err(AwaError::from)
103}
104
105async fn canonical_live_backlog(pool: &PgPool) -> Result<i64, AwaError> {
106    sqlx::query_scalar::<_, i64>(
107        r#"
108        SELECT
109            COALESCE((
110                SELECT count(*)::bigint
111                FROM awa.jobs_hot
112                WHERE state NOT IN ('completed', 'failed', 'cancelled')
113            ), 0)
114            +
115            COALESCE((
116                SELECT count(*)::bigint
117                FROM awa.scheduled_jobs
118            ), 0)
119        "#,
120    )
121    .fetch_one(pool)
122    .await
123    .map_err(AwaError::from)
124}
125
126pub async fn status_report(pool: &PgPool) -> Result<StorageStatusReport, AwaError> {
127    let status = status(pool).await?;
128    let canonical_live_backlog = canonical_live_backlog(pool).await?;
129    let prepared_queue_storage_schema = queue_storage_schema_from_status(&status);
130    let prepared_schema_ready = match prepared_queue_storage_schema.as_deref() {
131        Some(schema) => prepared_queue_storage_schema_ready(pool, schema).await?,
132        None => false,
133    };
134
135    let instances = admin::list_runtime_instances(pool).await?;
136    let mut live_runtime_capability_counts = BTreeMap::new();
137    for instance in instances.into_iter().filter(|instance| !instance.stale) {
138        *live_runtime_capability_counts
139            .entry(instance.storage_capability.as_str().to_string())
140            .or_insert(0) += 1;
141    }
142
143    let live_canonical = *live_runtime_capability_counts
144        .get("canonical")
145        .unwrap_or(&0);
146    let live_queue_storage = *live_runtime_capability_counts
147        .get("queue_storage")
148        .unwrap_or(&0);
149    let live_drain = *live_runtime_capability_counts
150        .get("canonical_drain_only")
151        .unwrap_or(&0);
152
153    let mut enter_mixed_transition_blockers = Vec::new();
154    if status.state != "prepared" {
155        enter_mixed_transition_blockers
156            .push(format!("state is '{}' (expected 'prepared')", status.state));
157    }
158    if status.prepared_engine.as_deref() != Some("queue_storage") {
159        enter_mixed_transition_blockers.push(format!(
160            "prepared_engine is {:?} (expected Some(\"queue_storage\"))",
161            status.prepared_engine
162        ));
163    }
164    if prepared_queue_storage_schema.is_some() && !prepared_schema_ready {
165        enter_mixed_transition_blockers
166            .push("prepared queue-storage schema is missing required tables".to_string());
167    }
168    if live_canonical > 0 {
169        enter_mixed_transition_blockers.push(format!(
170            "{live_canonical} canonical-only runtime(s) are still live"
171        ));
172    }
173    if live_queue_storage == 0 {
174        enter_mixed_transition_blockers
175            .push("no live queue-storage-capable runtimes are reporting".to_string());
176    }
177
178    let mut finalize_blockers = Vec::new();
179    if status.state != "mixed_transition" {
180        finalize_blockers.push(format!(
181            "state is '{}' (expected 'mixed_transition')",
182            status.state
183        ));
184    }
185    if status.prepared_engine.as_deref() != Some("queue_storage") {
186        finalize_blockers.push(format!(
187            "prepared_engine is {:?} (expected Some(\"queue_storage\"))",
188            status.prepared_engine
189        ));
190    }
191    if canonical_live_backlog > 0 {
192        finalize_blockers.push(format!(
193            "canonical live backlog is {canonical_live_backlog}"
194        ));
195    }
196    if live_canonical + live_drain > 0 {
197        finalize_blockers.push(format!(
198            "{} canonical or drain-only runtime(s) are still live",
199            live_canonical + live_drain
200        ));
201    }
202
203    Ok(StorageStatusReport {
204        status,
205        canonical_live_backlog,
206        prepared_queue_storage_schema,
207        prepared_schema_ready,
208        live_runtime_capability_counts,
209        can_enter_mixed_transition: enter_mixed_transition_blockers.is_empty(),
210        enter_mixed_transition_blockers,
211        can_finalize: finalize_blockers.is_empty(),
212        finalize_blockers,
213    })
214}
awa_model/storage.rs

awa_model/
storage.rs
