1use crate::BackupError;
4use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey, Checksum, S3ConditionalPut};
5use object_store::ObjectStore;
6use std::sync::Arc;
7use std::time::Duration;
8
9pub const DEFAULT_WAL_INTERVAL_MS: u64 = 1000;
11pub const DEFAULT_MANIFEST_INTERVAL_S: u64 = 60;
13pub const DEFAULT_RETAIN_N: usize = 7;
15pub const DEFAULT_LEASE_TTL_S: u64 = 30;
17pub const SHUTDOWN_MARGIN: Duration = Duration::from_secs(2);
19
20pub const DEFAULT_OP_TIMEOUT_MS: u64 = 10_000;
22pub const DEFAULT_SHUTDOWN_GRACE_MS: u64 = 10_000;
24
25#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct BackupConfig {
28 pub s3_url: String,
31 pub wal_interval_ms: u64,
33 pub manifest_interval_s: u64,
36 pub retain_n: usize,
38 pub sse: Option<String>,
41 pub lease_ttl_s: u64,
45 pub lease_conditional: bool,
50 pub op_timeout_ms: u64,
55 pub shutdown_grace_ms: u64,
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
62pub struct S3Location {
63 pub scheme: String,
65 pub bucket: String,
66 pub prefix: String,
69}
70
71impl BackupConfig {
72 pub fn from_env() -> Result<Option<Self>, BackupError> {
80 Self::from_lookup(|k| std::env::var(k).ok())
81 }
82
83 pub fn from_lookup(get: impl Fn(&str) -> Option<String>) -> Result<Option<Self>, BackupError> {
85 let Some(s3_url) = get("MUSHROOMDB_BACKUP_S3_URL").filter(|v| !v.trim().is_empty()) else {
86 return Ok(None);
87 };
88 let s3_url = s3_url.trim().to_string();
89 parse_location(&s3_url)?;
90 let number = |key: &str, default: u64| -> Result<u64, BackupError> {
91 match get(key).filter(|v| !v.trim().is_empty()) {
92 None => Ok(default),
93 Some(v) => v.trim().parse().map_err(|_| {
94 BackupError::Config(format!("{key} must be a whole number, got {v:?}"))
95 }),
96 }
97 };
98 let wal_interval_ms = number("MUSHROOMDB_BACKUP_WAL_INTERVAL_MS", DEFAULT_WAL_INTERVAL_MS)?;
99 if wal_interval_ms == 0 {
100 return Err(BackupError::Config(
101 "MUSHROOMDB_BACKUP_WAL_INTERVAL_MS must be at least 1".into(),
102 ));
103 }
104 Ok(Some(Self {
105 s3_url,
106 wal_interval_ms,
107 manifest_interval_s: number(
108 "MUSHROOMDB_BACKUP_MANIFEST_INTERVAL_S",
109 DEFAULT_MANIFEST_INTERVAL_S,
110 )?,
111 retain_n: number("MUSHROOMDB_BACKUP_RETAIN_N", DEFAULT_RETAIN_N as u64)? as usize,
112 sse: get("MUSHROOMDB_BACKUP_SSE")
113 .map(|v| v.trim().to_string())
114 .filter(|v| !v.is_empty()),
115 lease_ttl_s: positive(
116 number("MUSHROOMDB_BACKUP_LEASE_TTL", DEFAULT_LEASE_TTL_S)?,
117 "MUSHROOMDB_BACKUP_LEASE_TTL",
118 )?,
119 lease_conditional: match get("MUSHROOMDB_BACKUP_LEASE_CONDITIONAL")
120 .map(|v| v.trim().to_ascii_lowercase())
121 .filter(|v| !v.is_empty())
122 .as_deref()
123 {
124 None | Some("on") => true,
125 Some("off") => false,
126 Some(v) => {
127 return Err(BackupError::Config(format!(
128 "MUSHROOMDB_BACKUP_LEASE_CONDITIONAL must be on or off, got {v:?}"
129 )))
130 }
131 },
132 op_timeout_ms: positive(
133 number("MUSHROOMDB_BACKUP_OP_TIMEOUT_MS", DEFAULT_OP_TIMEOUT_MS)?,
134 "MUSHROOMDB_BACKUP_OP_TIMEOUT_MS",
135 )?,
136 shutdown_grace_ms: positive(
137 number(
138 "MUSHROOMDB_BACKUP_SHUTDOWN_GRACE_MS",
139 DEFAULT_SHUTDOWN_GRACE_MS,
140 )?,
141 "MUSHROOMDB_BACKUP_SHUTDOWN_GRACE_MS",
142 )?,
143 }))
144 }
145
146 pub fn new(s3_url: impl Into<String>) -> Self {
148 Self {
149 s3_url: s3_url.into(),
150 wal_interval_ms: DEFAULT_WAL_INTERVAL_MS,
151 manifest_interval_s: DEFAULT_MANIFEST_INTERVAL_S,
152 retain_n: DEFAULT_RETAIN_N,
153 sse: None,
154 lease_ttl_s: DEFAULT_LEASE_TTL_S,
155 lease_conditional: true,
156 op_timeout_ms: DEFAULT_OP_TIMEOUT_MS,
157 shutdown_grace_ms: DEFAULT_SHUTDOWN_GRACE_MS,
158 }
159 }
160
161 pub fn shutdown_budget(&self) -> Duration {
168 Duration::from_millis(self.shutdown_grace_ms.saturating_add(self.op_timeout_ms))
169 + SHUTDOWN_MARGIN
170 }
171
172 pub fn fit_shutdown_within(&mut self, within: Duration) {
179 let budget = self.shutdown_budget();
180 if budget <= within {
181 return;
182 }
183 let margin = SHUTDOWN_MARGIN.min(within / 5);
184 let room = within.saturating_sub(margin).as_secs_f64();
185 let used = (self.shutdown_grace_ms.saturating_add(self.op_timeout_ms)) as f64 / 1000.0;
186 let ratio = (room / used).min(1.0);
187 let scale = |ms: u64| ((ms as f64 * ratio) as u64).max(100);
188 self.shutdown_grace_ms = scale(self.shutdown_grace_ms);
189 self.op_timeout_ms = scale(self.op_timeout_ms);
190 }
191
192 pub fn location(&self) -> Result<S3Location, BackupError> {
194 parse_location(&self.s3_url)
195 }
196}
197
198fn positive(v: u64, key: &str) -> Result<u64, BackupError> {
199 if v == 0 {
200 return Err(BackupError::Config(format!("{key} must be at least 1")));
201 }
202 Ok(v)
203}
204
205pub fn parse_location(url: &str) -> Result<S3Location, BackupError> {
207 let bad = |why: &str| {
208 BackupError::Config(format!(
209 "MUSHROOMDB_BACKUP_S3_URL {url:?} {why}; expected s3://bucket/prefix, \
210 r2://bucket/prefix or minio://bucket/prefix"
211 ))
212 };
213 let (scheme, rest) = url.split_once("://").ok_or_else(|| bad("has no scheme"))?;
214 if !matches!(scheme, "s3" | "r2" | "minio") {
215 return Err(bad("has an unsupported scheme"));
216 }
217 let (bucket, prefix) = rest.split_once('/').unwrap_or((rest, ""));
218 if bucket.is_empty() {
219 return Err(bad("names no bucket"));
220 }
221 Ok(S3Location {
222 scheme: scheme.to_string(),
223 bucket: bucket.to_string(),
224 prefix: prefix.trim_matches('/').to_string(),
225 })
226}
227
228pub fn build_store(config: &BackupConfig) -> Result<Arc<dyn ObjectStore>, BackupError> {
239 build_store_with(config, |k| std::env::var(k).ok())
240}
241
242pub fn build_store_with(
244 config: &BackupConfig,
245 get: impl Fn(&str) -> Option<String>,
246) -> Result<Arc<dyn ObjectStore>, BackupError> {
247 let loc = config.location()?;
248 let get = |k: &str| {
249 get(k)
250 .map(|v| v.trim().to_string())
251 .filter(|v| !v.is_empty())
252 };
253 let endpoint = get("AWS_ENDPOINT_URL");
254 if loc.scheme != "s3" && endpoint.is_none() {
255 return Err(BackupError::Config(format!(
256 "{}:// needs AWS_ENDPOINT_URL set to the service's endpoint",
257 loc.scheme
258 )));
259 }
260 let mut b = AmazonS3Builder::new()
261 .with_bucket_name(&loc.bucket)
262 .with_region(get("AWS_REGION").unwrap_or_else(|| "us-east-1".into()))
263 .with_checksum_algorithm(Checksum::SHA256);
264 if config.lease_conditional {
265 b = b.with_conditional_put(S3ConditionalPut::ETagMatch);
266 }
267 if let Some(id) = get("AWS_ACCESS_KEY_ID") {
268 b = b.with_access_key_id(id);
269 }
270 if let Some(secret) = get("AWS_SECRET_ACCESS_KEY") {
271 b = b.with_secret_access_key(secret);
272 }
273 if let Some(token) = get("AWS_SESSION_TOKEN") {
274 b = b.with_token(token);
275 }
276 if let Some(ep) = endpoint {
277 b = b
278 .with_allow_http(ep.starts_with("http://"))
279 .with_endpoint(ep);
280 }
281 if let Some(sse) = &config.sse {
282 let key: AmazonS3ConfigKey = "aws_server_side_encryption"
283 .parse()
284 .map_err(|e| BackupError::Config(format!("{e}")))?;
285 b = b.with_config(key, sse);
286 }
287 let store = b
288 .build()
289 .map_err(|e| BackupError::Config(format!("cannot build the S3 client: {e}")))?;
290 Ok(Arc::new(store))
291}
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296 use std::collections::HashMap;
297
298 fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
299 let m: HashMap<String, String> = pairs
300 .iter()
301 .map(|(k, v)| (k.to_string(), v.to_string()))
302 .collect();
303 move |k| m.get(k).cloned()
304 }
305
306 #[test]
307 fn unset_url_means_backup_is_off() {
308 assert_eq!(BackupConfig::from_lookup(env(&[])).unwrap(), None);
309 assert_eq!(
310 BackupConfig::from_lookup(env(&[("MUSHROOMDB_BACKUP_S3_URL", " ")])).unwrap(),
311 None
312 );
313 }
314
315 #[test]
316 fn defaults_apply_when_only_the_url_is_set() {
317 let c = BackupConfig::from_lookup(env(&[("MUSHROOMDB_BACKUP_S3_URL", "s3://b/p")]))
318 .unwrap()
319 .unwrap();
320 assert_eq!(c, BackupConfig::new("s3://b/p"));
321 assert_eq!(c.wal_interval_ms, 1000);
322 assert_eq!(c.manifest_interval_s, 60);
323 }
324
325 #[test]
326 fn every_setting_is_read() {
327 let c = BackupConfig::from_lookup(env(&[
328 ("MUSHROOMDB_BACKUP_S3_URL", "minio://bk/a/b/"),
329 ("MUSHROOMDB_BACKUP_WAL_INTERVAL_MS", "250"),
330 ("MUSHROOMDB_BACKUP_MANIFEST_INTERVAL_S", "5"),
331 ("MUSHROOMDB_BACKUP_RETAIN_N", "0"),
332 ("MUSHROOMDB_BACKUP_SSE", "AES256"),
333 ("MUSHROOMDB_BACKUP_LEASE_TTL", "9"),
334 ("MUSHROOMDB_BACKUP_OP_TIMEOUT_MS", "1500"),
335 ("MUSHROOMDB_BACKUP_SHUTDOWN_GRACE_MS", "2500"),
336 ("MUSHROOMDB_BACKUP_LEASE_CONDITIONAL", "off"),
337 ]))
338 .unwrap()
339 .unwrap();
340 assert_eq!(
341 (
342 c.wal_interval_ms,
343 c.manifest_interval_s,
344 c.retain_n,
345 c.lease_ttl_s
346 ),
347 (250, 5, 0, 9)
348 );
349 assert_eq!(c.sse.as_deref(), Some("AES256"));
350 assert_eq!((c.op_timeout_ms, c.shutdown_grace_ms), (1500, 2500));
351 assert!(!c.lease_conditional);
352 let loc = c.location().unwrap();
353 assert_eq!((loc.bucket.as_str(), loc.prefix.as_str()), ("bk", "a/b"));
354 }
355
356 #[test]
357 fn malformed_values_are_errors_not_silence() {
358 for (k, v) in [
359 ("MUSHROOMDB_BACKUP_S3_URL", "gs://b/p"),
360 ("MUSHROOMDB_BACKUP_S3_URL", "bucket/p"),
361 ("MUSHROOMDB_BACKUP_S3_URL", "s3:///p"),
362 ] {
363 assert!(BackupConfig::from_lookup(env(&[(k, v)])).is_err(), "{v}");
364 }
365 let bad_n = env(&[
366 ("MUSHROOMDB_BACKUP_S3_URL", "s3://b"),
367 ("MUSHROOMDB_BACKUP_WAL_INTERVAL_MS", "fast"),
368 ]);
369 assert!(BackupConfig::from_lookup(bad_n).is_err());
370 let zero = env(&[
371 ("MUSHROOMDB_BACKUP_S3_URL", "s3://b"),
372 ("MUSHROOMDB_BACKUP_WAL_INTERVAL_MS", "0"),
373 ]);
374 assert!(BackupConfig::from_lookup(zero).is_err());
375 for (k, v) in [
376 ("MUSHROOMDB_BACKUP_LEASE_TTL", "0"),
377 ("MUSHROOMDB_BACKUP_LEASE_CONDITIONAL", "maybe"),
378 ] {
379 let e = env(&[("MUSHROOMDB_BACKUP_S3_URL", "s3://b"), (k, v)]);
380 assert!(BackupConfig::from_lookup(e).is_err(), "{k}={v}");
381 }
382 let on = env(&[
383 ("MUSHROOMDB_BACKUP_S3_URL", "s3://b"),
384 ("MUSHROOMDB_BACKUP_LEASE_CONDITIONAL", " ON "),
385 ]);
386 assert!(
387 BackupConfig::from_lookup(on)
388 .unwrap()
389 .unwrap()
390 .lease_conditional
391 );
392 }
393
394 #[test]
395 fn the_shutdown_budget_covers_the_flush_and_the_lease_release() {
396 let mut c = BackupConfig::new("s3://b/p");
397 assert_eq!(c.shutdown_budget(), Duration::from_secs(22));
398 c.shutdown_grace_ms = 3_000;
399 c.op_timeout_ms = 1_000;
400 assert_eq!(c.shutdown_budget(), Duration::from_secs(6));
401 }
402
403 #[test]
404 fn a_tight_shutdown_window_shrinks_the_bounds_to_fit() {
405 let mut c = BackupConfig::new("s3://b/p");
406 c.fit_shutdown_within(Duration::from_secs(30));
407 assert_eq!(c.shutdown_budget(), Duration::from_secs(22), "already fits");
408 c.fit_shutdown_within(Duration::from_secs(10));
409 let worst = |c: &BackupConfig, within: Duration| {
410 Duration::from_millis(c.shutdown_grace_ms + c.op_timeout_ms)
411 + SHUTDOWN_MARGIN.min(within / 5)
412 };
413 assert!(
414 worst(&c, Duration::from_secs(10)) <= Duration::from_secs(10),
415 "{c:?}"
416 );
417 assert!(c.op_timeout_ms >= 100 && c.shutdown_grace_ms >= 100);
418 c.fit_shutdown_within(Duration::from_millis(300));
419 assert!(
420 worst(&c, Duration::from_millis(300)) <= Duration::from_millis(300),
421 "{c:?}"
422 );
423 }
424
425 #[test]
426 fn a_non_aws_scheme_needs_an_endpoint() {
427 let c = BackupConfig::new("minio://b/p");
428 assert!(build_store_with(&c, env(&[])).is_err());
429 assert!(
430 build_store_with(&c, env(&[("AWS_ENDPOINT_URL", "http://127.0.0.1:9000")])).is_ok()
431 );
432 assert!(build_store_with(&BackupConfig::new("s3://b/p"), env(&[])).is_ok());
433 }
434}