alopex_server/ops/
memory.rs

1use std::env;
2use std::path::PathBuf;
3use std::sync::Arc;
4
5use alopex_core::sql::stream::DEFAULT_SPILL_THRESHOLD_BYTES;
6use alopex_sql::executor::memory::{MemoryPolicy, SpillMetricsSink, SpillPolicy as SqlSpillPolicy};
7
8use crate::error::{Result, ServerError};
9use crate::metrics::Metrics;
10
11#[derive(Clone, Debug)]
12pub enum SpillPolicy {
13    FailFast,
14    SpillToDisk { directory: PathBuf },
15}
16
17#[derive(Clone)]
18pub struct MemoryControlPolicy {
19    limit_bytes: Option<u64>,
20    spill_policy: SpillPolicy,
21    metrics: Option<Metrics>,
22}
23
24impl MemoryControlPolicy {
25    pub fn from_env() -> Self {
26        let limit_bytes = env::var("ALOPEX_MEMORY_LIMIT_BYTES")
27            .ok()
28            .and_then(|val| val.parse::<u64>().ok())
29            .unwrap_or(DEFAULT_SPILL_THRESHOLD_BYTES);
30
31        let policy = env::var("ALOPEX_MEMORY_SPILL_POLICY")
32            .unwrap_or_else(|_| "fail_fast".to_string())
33            .to_ascii_lowercase();
34        let spill_dir = env::var("ALOPEX_MEMORY_SPILL_DIR").ok().map(PathBuf::from);
35
36        let spill_policy = match policy.as_str() {
37            "spill" | "spill_to_disk" | "spill-to-disk" => spill_dir
38                .map(|directory| SpillPolicy::SpillToDisk { directory })
39                .unwrap_or(SpillPolicy::FailFast),
40            _ => SpillPolicy::FailFast,
41        };
42
43        Self {
44            limit_bytes: Some(limit_bytes),
45            spill_policy,
46            metrics: None,
47        }
48    }
49
50    pub fn from_env_with_metrics(metrics: Metrics) -> Self {
51        Self::from_env().with_metrics(metrics)
52    }
53
54    pub fn with_metrics(mut self, metrics: Metrics) -> Self {
55        self.metrics = Some(metrics);
56        self
57    }
58
59    pub fn limit_bytes(&self) -> Option<u64> {
60        self.limit_bytes
61    }
62
63    pub fn spill_policy(&self) -> &SpillPolicy {
64        &self.spill_policy
65    }
66
67    pub fn sql_policy(&self) -> Option<MemoryPolicy> {
68        let limit = self.limit_bytes?;
69        let spill_policy = match &self.spill_policy {
70            SpillPolicy::FailFast => SqlSpillPolicy::FailFast,
71            SpillPolicy::SpillToDisk { directory } => SqlSpillPolicy::SpillToDisk {
72                directory: directory.clone(),
73            },
74        };
75        let mut policy = MemoryPolicy::new(Some(limit), spill_policy);
76        if let Some(metrics) = &self.metrics {
77            policy = policy.with_metrics(Arc::new(MetricsSpillSink {
78                metrics: metrics.clone(),
79            }));
80        }
81        Some(policy)
82    }
83
84    pub fn enforce_output_bytes(&self, bytes: u64) -> Result<()> {
85        let Some(limit) = self.limit_bytes else {
86            return Ok(());
87        };
88        if bytes <= limit {
89            return Ok(());
90        }
91        self.enforce_limit(limit, bytes)
92    }
93
94    fn enforce_limit(&self, limit: u64, bytes: u64) -> Result<()> {
95        match &self.spill_policy {
96            SpillPolicy::FailFast => Err(ServerError::PayloadTooLarge(format!(
97                "query memory limit exceeded: {bytes} bytes (limit {limit})"
98            ))),
99            SpillPolicy::SpillToDisk { .. } => Err(ServerError::PayloadTooLarge(format!(
100                "query output exceeds memory limit: {bytes} bytes (limit {limit})"
101            ))),
102        }
103    }
104}
105
106impl std::fmt::Debug for MemoryControlPolicy {
107    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        f.debug_struct("MemoryControlPolicy")
109            .field("limit_bytes", &self.limit_bytes)
110            .field("spill_policy", &self.spill_policy)
111            .finish()
112    }
113}
114
115struct MetricsSpillSink {
116    metrics: Metrics,
117}
118
119impl SpillMetricsSink for MetricsSpillSink {
120    fn record_spill(&self, bytes: u64, files: u64) {
121        self.metrics.record_spill(bytes, files);
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use alopex_core::sql::stream::DEFAULT_SPILL_THRESHOLD_BYTES;
129    use std::sync::{Mutex, MutexGuard};
130
131    static ENV_LOCK: Mutex<()> = Mutex::new(());
132
133    struct EnvVarGuard {
134        key: &'static str,
135        value: Option<String>,
136        _lock: MutexGuard<'static, ()>,
137    }
138
139    impl EnvVarGuard {
140        fn unset(key: &'static str) -> Self {
141            let lock = ENV_LOCK.lock().unwrap();
142            let value = env::var(key).ok();
143            // SAFETY: This test module serializes all mutations of this env var with ENV_LOCK.
144            unsafe {
145                env::remove_var(key);
146            }
147            Self {
148                key,
149                value,
150                _lock: lock,
151            }
152        }
153    }
154
155    impl Drop for EnvVarGuard {
156        fn drop(&mut self) {
157            // SAFETY: The guard still holds ENV_LOCK, so restoration is serialized.
158            unsafe {
159                if let Some(value) = &self.value {
160                    env::set_var(self.key, value);
161                } else {
162                    env::remove_var(self.key);
163                }
164            }
165        }
166    }
167
168    #[test]
169    fn from_env_uses_default_spill_threshold_when_limit_is_unset() {
170        let _guard = EnvVarGuard::unset("ALOPEX_MEMORY_LIMIT_BYTES");
171
172        let policy = MemoryControlPolicy::from_env();
173
174        assert_eq!(policy.limit_bytes(), Some(DEFAULT_SPILL_THRESHOLD_BYTES));
175    }
176}
alopex_server/ops/memory.rs

alopex_server/ops/
memory.rs
