Constellation, Spacedust, Slingshot, UFOs: atproto crates and services for microcosm
background task for backups
+134
-33
+25
-15
constellation/src/bin/main.rs
+25
-15
constellation/src/bin/main.rs
···
1
-
use anyhow::Result;
1
+
use anyhow::{bail, Result};
2
2
use clap::{Parser, ValueEnum};
3
3
use metrics_exporter_prometheus::PrometheusBuilder;
4
4
use std::num::NonZero;
···
37
37
/// Initiate a database backup into this dir, if supported by the storage
38
38
#[arg(long)]
39
39
backup: Option<PathBuf>,
40
+
/// Start a background task to take backups every N hours
41
+
#[arg(long)]
42
+
backup_interval: Option<u64>,
43
+
/// If backup_interval is configured, purge the oldest backup once this many backups are saved
44
+
#[arg(long)]
45
+
max_old_backups: Option<usize>,
40
46
/// Saved jsonl from jetstream to use instead of a live subscription
41
47
#[arg(short, long)]
42
48
fixture: Option<PathBuf>,
···
71
77
72
78
let stream = jetstream_url(&args.jetstream);
73
79
println!("using jetstream server {stream:?}...",);
80
+
81
+
let stay_alive = CancellationToken::new();
74
82
75
83
match args.backend {
76
-
StorageBackend::Memory => run(MemStorage::new(), fixture, None, stream),
84
+
StorageBackend::Memory => run(MemStorage::new(), fixture, None, stream, stay_alive),
77
85
#[cfg(feature = "rocks")]
78
86
StorageBackend::Rocks => {
79
87
let storage_dir = args.data.clone().unwrap_or("rocks.test".into());
80
88
println!("starting rocksdb...");
81
-
let rocks = RocksStorage::new(storage_dir)?;
89
+
let mut rocks = RocksStorage::new(storage_dir)?;
82
90
if let Some(backup_dir) = args.backup {
83
-
rocks.start_backup(backup_dir)?;
91
+
let auto_backup = match (args.backup_interval, args.max_old_backups) {
92
+
(Some(interval_hrs), copies) => Some((interval_hrs, copies)),
93
+
(None, None) => None,
94
+
(None, Some(_)) => bail!("invalid backup config: --max-old-backups requires --backup-interval to be configured"),
95
+
};
96
+
rocks.start_backup(backup_dir, auto_backup, stay_alive.clone())?;
84
97
}
85
98
println!("rocks ready.");
86
-
run(rocks, fixture, args.data, stream)
99
+
run(rocks, fixture, args.data, stream, stay_alive)
87
100
}
88
101
}
89
102
}
···
93
106
fixture: Option<PathBuf>,
94
107
data_dir: Option<PathBuf>,
95
108
stream: String,
109
+
stay_alive: CancellationToken,
96
110
) -> Result<()> {
97
-
let stay_alive = CancellationToken::new();
98
-
99
111
ctrlc::set_handler({
100
112
let mut desperation: u8 = 0;
101
113
let stay_alive = stay_alive.clone();
102
-
move || {
103
-
match desperation {
104
-
0 => {
105
-
println!("ok, shutting down...");
106
-
stay_alive.cancel();
107
-
}
108
-
1.. => panic!("fine, panicking!"),
114
+
move || match desperation {
115
+
0 => {
116
+
println!("ok, shutting down...");
117
+
stay_alive.cancel();
118
+
desperation += 1;
109
119
}
110
-
desperation += 1;
120
+
1.. => panic!("fine, panicking!"),
111
121
}
112
122
})?;
113
123
+109
-18
constellation/src/storage/rocks_store.rs
+109
-18
constellation/src/storage/rocks_store.rs
···
4
4
use bincode::Options as BincodeOptions;
5
5
use links::CollectedLink;
6
6
use metrics::{counter, describe_counter, describe_histogram, histogram, Unit};
7
+
use ratelimit::Ratelimiter;
8
+
use rocksdb::backup::{BackupEngine, BackupEngineOptions};
7
9
use rocksdb::{
8
10
AsColumnFamilyRef, ColumnFamilyDescriptor, DBWithThreadMode, IteratorMode, MergeOperands,
9
11
MultiThreaded, Options, PrefixRange, ReadOptions, WriteBatch,
···
12
14
use std::collections::{HashMap, HashSet};
13
15
use std::io::Read;
14
16
use std::marker::PhantomData;
15
-
use std::path::Path;
17
+
use std::path::{Path, PathBuf};
16
18
use std::sync::{
17
19
atomic::{AtomicU64, Ordering},
18
20
Arc,
19
21
};
20
-
use std::time::Instant;
22
+
use std::thread;
23
+
use std::time::{Duration, Instant};
24
+
use tokio_util::sync::CancellationToken;
21
25
22
26
static DID_IDS_CF: &str = "did_ids";
23
27
static TARGET_IDS_CF: &str = "target_ids";
···
55
59
did_id_table: IdTable<Did, DidIdValue, true>,
56
60
target_id_table: IdTable<TargetKey, TargetId, false>,
57
61
is_writer: bool,
62
+
backup_task: Arc<Option<thread::JoinHandle<Result<()>>>>,
58
63
}
59
64
60
65
trait IdTableValue: ValueFromRocks + Clone {
···
289
294
did_id_table,
290
295
target_id_table,
291
296
is_writer: true,
297
+
backup_task: None.into(),
292
298
})
293
299
}
294
300
295
-
pub fn start_backup(&self, path: impl AsRef<Path>) -> Result<()> {
296
-
use rocksdb::backup::{BackupEngine, BackupEngineOptions};
297
-
eprintln!("getting ready to start backup...");
298
-
let mut engine =
299
-
BackupEngine::open(&BackupEngineOptions::new(path)?, &rocksdb::Env::new()?)?;
300
-
std::thread::spawn({
301
-
let db = self.db.clone();
302
-
move || {
303
-
eprintln!("backup starting.");
304
-
let t0 = Instant::now();
305
-
if let Err(e) = engine.create_new_backup(&db) {
306
-
eprintln!("oh no, backup failed: {e:?}");
301
+
pub fn start_backup(
302
+
&mut self,
303
+
path: PathBuf,
304
+
auto: Option<(u64, Option<usize>)>,
305
+
stay_alive: CancellationToken,
306
+
) -> Result<()> {
307
+
let task = if let Some((interval_hrs, copies)) = auto {
308
+
eprintln!("backups: starting background task...");
309
+
self.backup_task(path, interval_hrs, copies, stay_alive)
310
+
} else {
311
+
eprintln!("backups: starting a one-off backup...");
312
+
thread::spawn({
313
+
let db = self.db.clone();
314
+
move || Self::do_backup(db, path)
315
+
})
316
+
};
317
+
self.backup_task = Arc::new(Some(task));
318
+
Ok(())
319
+
}
320
+
321
+
fn backup_task(
322
+
&self,
323
+
path: PathBuf,
324
+
interval_hrs: u64,
325
+
copies: Option<usize>,
326
+
stay_alive: CancellationToken,
327
+
) -> std::thread::JoinHandle<Result<()>> {
328
+
let db = self.db.clone();
329
+
thread::spawn(move || {
330
+
let limit =
331
+
Ratelimiter::builder(1, Duration::from_secs(interval_hrs * 60 * 60)).build()?;
332
+
let minimum_sleep = Duration::from_secs(1);
333
+
334
+
'quit: loop {
335
+
if let Err(sleep) = limit.try_wait() {
336
+
eprintln!("backups: background: next backup scheduled in {sleep:?}");
337
+
let waiting = Instant::now();
338
+
loop {
339
+
let remaining = sleep - waiting.elapsed();
340
+
if stay_alive.is_cancelled() {
341
+
break 'quit;
342
+
} else if remaining <= Duration::ZERO {
343
+
break;
344
+
} else if remaining < minimum_sleep {
345
+
thread::sleep(remaining);
346
+
break;
347
+
} else {
348
+
thread::sleep(minimum_sleep);
349
+
}
350
+
}
351
+
}
352
+
eprintln!("backups: background: starting backup...");
353
+
if let Err(e) = Self::do_backup(db.clone(), &path) {
354
+
eprintln!("backups: background: backup failed: {e:?}");
355
+
// todo: metrics
307
356
} else {
308
-
eprintln!("yay, backup worked?");
357
+
eprintln!("backups: background: backup succeeded yay");
358
+
}
359
+
if let Some(copies) = copies {
360
+
eprintln!("backups: background: trimming to {copies} saved backups...");
361
+
if let Err(e) = Self::trim_backups(copies, &path) {
362
+
eprintln!("backups: background: failed to trim backups: {e:?}");
363
+
} else {
364
+
eprintln!("backups: background: trimming worked!")
365
+
}
309
366
}
310
-
eprintln!("backup finished after {:.2}s", t0.elapsed().as_secs_f32());
311
367
}
312
-
});
313
-
eprintln!("backups should be happening in bg thread.");
368
+
369
+
Ok(())
370
+
})
371
+
}
372
+
373
+
fn do_backup(db: Arc<DBWithThreadMode<MultiThreaded>>, path: impl AsRef<Path>) -> Result<()> {
374
+
let mut engine =
375
+
BackupEngine::open(&BackupEngineOptions::new(path)?, &rocksdb::Env::new()?)?;
376
+
eprintln!("backups: starting backup...");
377
+
let t0 = Instant::now();
378
+
if let Err(e) = engine.create_new_backup(&db) {
379
+
eprintln!("backups: oh no, backup failed: {e:?}");
380
+
} else {
381
+
eprintln!("backups: yay, backup worked?");
382
+
}
383
+
eprintln!(
384
+
"backups: backup finished after {:.2}s",
385
+
t0.elapsed().as_secs_f32()
386
+
);
387
+
Ok(())
388
+
}
389
+
390
+
fn trim_backups(num_backups_to_keep: usize, path: impl AsRef<Path>) -> Result<()> {
391
+
let mut engine =
392
+
BackupEngine::open(&BackupEngineOptions::new(path)?, &rocksdb::Env::new()?)?;
393
+
engine.purge_old_backups(num_backups_to_keep)?;
314
394
Ok(())
315
395
}
316
396
···
646
726
opt
647
727
}) {
648
728
eprintln!("rocks: flushing memtables failed: {e:?}");
729
+
}
730
+
match Arc::get_mut(&mut self.backup_task) {
731
+
Some(maybe_task) => {
732
+
if let Some(task) = maybe_task.take() {
733
+
eprintln!("waiting for backup task to complete...");
734
+
if let Err(e) = task.join() {
735
+
eprintln!("failed to join backup task: {e:?}");
736
+
}
737
+
}
738
+
}
739
+
None => eprintln!("rocks: failed to get backup task, likely a bug."),
649
740
}
650
741
self.db.cancel_all_background_work(true);
651
742
}