commit 8f368ff7cbe4874d24465ce52f8a36e0e814a463 · nekomimi.pet/microcosm-rs

+25 -15

constellation/src/bin/main.rs

···

       1
       1
       -
       use anyhow::Result;

     

       1
       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
       40
       +
           /// Start a background task to take backups every N hours

     

       41
       41
       +
           #[arg(long)]

     

       42
       42
       +
           backup_interval: Option<u64>,

     

       43
       43
       +
           /// If backup_interval is configured, purge the oldest backup once this many backups are saved

     

       44
       44
       +
           #[arg(long)]

     

       45
       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
       80
       +
       

     

       81
       81
       +
           let stay_alive = CancellationToken::new();

     

       74
       82
        
       

     

       75
       83
        
           match args.backend {

     

       76
       76
       -
               StorageBackend::Memory => run(MemStorage::new(), fixture, None, stream),

     

       84
       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
       81
       -
                   let rocks = RocksStorage::new(storage_dir)?;

     

       89
       89
       +
                   let mut rocks = RocksStorage::new(storage_dir)?;

     

       82
       90
        
                   if let Some(backup_dir) = args.backup {

     

       83
       83
       -
                       rocks.start_backup(backup_dir)?;

     

       91
       91
       +
                       let auto_backup = match (args.backup_interval, args.max_old_backups) {

     

       92
       92
       +
                           (Some(interval_hrs), copies) => Some((interval_hrs, copies)),

     

       93
       93
       +
                           (None, None) => None,

     

       94
       94
       +
                           (None, Some(_)) => bail!("invalid backup config: --max-old-backups requires --backup-interval to be configured"),

     

       95
       95
       +
                       };

     

       96
       96
       +
                       rocks.start_backup(backup_dir, auto_backup, stay_alive.clone())?;

     

       84
       97
        
                   }

     

       85
       98
        
                   println!("rocks ready.");

     

       86
       86
       -
                   run(rocks, fixture, args.data, stream)

     

       99
       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
       109
       +
           stay_alive: CancellationToken,

     

       96
       110
        
       ) -> Result<()> {

     

       97
       97
       -
           let stay_alive = CancellationToken::new();

     

       98
       98
       -
       

     

       99
       111
        
           ctrlc::set_handler({

     

       100
       112
        
               let mut desperation: u8 = 0;

     

       101
       113
        
               let stay_alive = stay_alive.clone();

     

       102
       102
       -
               move || {

     

       103
       103
       -
                   match desperation {

     

       104
       104
       -
                       0 => {

     

       105
       105
       -
                           println!("ok, shutting down...");

     

       106
       106
       -
                           stay_alive.cancel();

     

       107
       107
       -
                       }

     

       108
       108
       -
                       1.. => panic!("fine, panicking!"),

     

       114
       114
       +
               move || match desperation {

     

       115
       115
       +
                   0 => {

     

       116
       116
       +
                       println!("ok, shutting down...");

     

       117
       117
       +
                       stay_alive.cancel();

     

       118
       118
       +
                       desperation += 1;

     

       109
       119
        
                   }

     

       110
       110
       -
                   desperation += 1;

     

       120
       120
       +
                   1.. => panic!("fine, panicking!"),

     

       111
       121
        
               }

     

       112
       122
        
           })?;

     

       113
       123

+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
       7
       +
       use ratelimit::Ratelimiter;

     

       8
       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
       15
       -
       use std::path::Path;

     

       17
       17
       +
       use std::path::{Path, PathBuf};

     

       16
       18
        
       use std::sync::{

     

       17
       19
        
           atomic::{AtomicU64, Ordering},

     

       18
       20
        
           Arc,

     

       19
       21
        
       };

     

       20
       20
       -
       use std::time::Instant;

     

       22
       22
       +
       use std::thread;

     

       23
       23
       +
       use std::time::{Duration, Instant};

     

       24
       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
       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
       297
       +
                   backup_task: None.into(),

     

       292
       298
        
               })

     

       293
       299
        
           }

     

       294
       300
        
       

     

       295
       295
       -
           pub fn start_backup(&self, path: impl AsRef<Path>) -> Result<()> {

     

       296
       296
       -
               use rocksdb::backup::{BackupEngine, BackupEngineOptions};

     

       297
       297
       -
               eprintln!("getting ready to start backup...");

     

       298
       298
       -
               let mut engine =

     

       299
       299
       -
                   BackupEngine::open(&BackupEngineOptions::new(path)?, &rocksdb::Env::new()?)?;

     

       300
       300
       -
               std::thread::spawn({

     

       301
       301
       -
                   let db = self.db.clone();

     

       302
       302
       -
                   move || {

     

       303
       303
       -
                       eprintln!("backup starting.");

     

       304
       304
       -
                       let t0 = Instant::now();

     

       305
       305
       -
                       if let Err(e) = engine.create_new_backup(&db) {

     

       306
       306
       -
                           eprintln!("oh no, backup failed: {e:?}");

     

       301
       301
       +
           pub fn start_backup(

     

       302
       302
       +
               &mut self,

     

       303
       303
       +
               path: PathBuf,

     

       304
       304
       +
               auto: Option<(u64, Option<usize>)>,

     

       305
       305
       +
               stay_alive: CancellationToken,

     

       306
       306
       +
           ) -> Result<()> {

     

       307
       307
       +
               let task = if let Some((interval_hrs, copies)) = auto {

     

       308
       308
       +
                   eprintln!("backups: starting background task...");

     

       309
       309
       +
                   self.backup_task(path, interval_hrs, copies, stay_alive)

     

       310
       310
       +
               } else {

     

       311
       311
       +
                   eprintln!("backups: starting a one-off backup...");

     

       312
       312
       +
                   thread::spawn({

     

       313
       313
       +
                       let db = self.db.clone();

     

       314
       314
       +
                       move || Self::do_backup(db, path)

     

       315
       315
       +
                   })

     

       316
       316
       +
               };

     

       317
       317
       +
               self.backup_task = Arc::new(Some(task));

     

       318
       318
       +
               Ok(())

     

       319
       319
       +
           }

     

       320
       320
       +
       

     

       321
       321
       +
           fn backup_task(

     

       322
       322
       +
               &self,

     

       323
       323
       +
               path: PathBuf,

     

       324
       324
       +
               interval_hrs: u64,

     

       325
       325
       +
               copies: Option<usize>,

     

       326
       326
       +
               stay_alive: CancellationToken,

     

       327
       327
       +
           ) -> std::thread::JoinHandle<Result<()>> {

     

       328
       328
       +
               let db = self.db.clone();

     

       329
       329
       +
               thread::spawn(move || {

     

       330
       330
       +
                   let limit =

     

       331
       331
       +
                       Ratelimiter::builder(1, Duration::from_secs(interval_hrs * 60 * 60)).build()?;

     

       332
       332
       +
                   let minimum_sleep = Duration::from_secs(1);

     

       333
       333
       +
       

     

       334
       334
       +
                   'quit: loop {

     

       335
       335
       +
                       if let Err(sleep) = limit.try_wait() {

     

       336
       336
       +
                           eprintln!("backups: background: next backup scheduled in {sleep:?}");

     

       337
       337
       +
                           let waiting = Instant::now();

     

       338
       338
       +
                           loop {

     

       339
       339
       +
                               let remaining = sleep - waiting.elapsed();

     

       340
       340
       +
                               if stay_alive.is_cancelled() {

     

       341
       341
       +
                                   break 'quit;

     

       342
       342
       +
                               } else if remaining <= Duration::ZERO {

     

       343
       343
       +
                                   break;

     

       344
       344
       +
                               } else if remaining < minimum_sleep {

     

       345
       345
       +
                                   thread::sleep(remaining);

     

       346
       346
       +
                                   break;

     

       347
       347
       +
                               } else {

     

       348
       348
       +
                                   thread::sleep(minimum_sleep);

     

       349
       349
       +
                               }

     

       350
       350
       +
                           }

     

       351
       351
       +
                       }

     

       352
       352
       +
                       eprintln!("backups: background: starting backup...");

     

       353
       353
       +
                       if let Err(e) = Self::do_backup(db.clone(), &path) {

     

       354
       354
       +
                           eprintln!("backups: background: backup failed: {e:?}");

     

       355
       355
       +
                           // todo: metrics

     

       307
       356
        
                       } else {

     

       308
       308
       -
                           eprintln!("yay, backup worked?");

     

       357
       357
       +
                           eprintln!("backups: background: backup succeeded yay");

     

       358
       358
       +
                       }

     

       359
       359
       +
                       if let Some(copies) = copies {

     

       360
       360
       +
                           eprintln!("backups: background: trimming to {copies} saved backups...");

     

       361
       361
       +
                           if let Err(e) = Self::trim_backups(copies, &path) {

     

       362
       362
       +
                               eprintln!("backups: background: failed to trim backups: {e:?}");

     

       363
       363
       +
                           } else {

     

       364
       364
       +
                               eprintln!("backups: background: trimming worked!")

     

       365
       365
       +
                           }

     

       309
       366
        
                       }

     

       310
       310
       -
                       eprintln!("backup finished after {:.2}s", t0.elapsed().as_secs_f32());

     

       311
       367
        
                   }

     

       312
       312
       -
               });

     

       313
       313
       -
               eprintln!("backups should be happening in bg thread.");

     

       368
       368
       +
       

     

       369
       369
       +
                   Ok(())

     

       370
       370
       +
               })

     

       371
       371
       +
           }

     

       372
       372
       +
       

     

       373
       373
       +
           fn do_backup(db: Arc<DBWithThreadMode<MultiThreaded>>, path: impl AsRef<Path>) -> Result<()> {

     

       374
       374
       +
               let mut engine =

     

       375
       375
       +
                   BackupEngine::open(&BackupEngineOptions::new(path)?, &rocksdb::Env::new()?)?;

     

       376
       376
       +
               eprintln!("backups: starting backup...");

     

       377
       377
       +
               let t0 = Instant::now();

     

       378
       378
       +
               if let Err(e) = engine.create_new_backup(&db) {

     

       379
       379
       +
                   eprintln!("backups: oh no, backup failed: {e:?}");

     

       380
       380
       +
               } else {

     

       381
       381
       +
                   eprintln!("backups: yay, backup worked?");

     

       382
       382
       +
               }

     

       383
       383
       +
               eprintln!(

     

       384
       384
       +
                   "backups: backup finished after {:.2}s",

     

       385
       385
       +
                   t0.elapsed().as_secs_f32()

     

       386
       386
       +
               );

     

       387
       387
       +
               Ok(())

     

       388
       388
       +
           }

     

       389
       389
       +
       

     

       390
       390
       +
           fn trim_backups(num_backups_to_keep: usize, path: impl AsRef<Path>) -> Result<()> {

     

       391
       391
       +
               let mut engine =

     

       392
       392
       +
                   BackupEngine::open(&BackupEngineOptions::new(path)?, &rocksdb::Env::new()?)?;

     

       393
       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
       729
       +
                   }

     

       730
       730
       +
                   match Arc::get_mut(&mut self.backup_task) {

     

       731
       731
       +
                       Some(maybe_task) => {

     

       732
       732
       +
                           if let Some(task) = maybe_task.take() {

     

       733
       733
       +
                               eprintln!("waiting for backup task to complete...");

     

       734
       734
       +
                               if let Err(e) = task.join() {

     

       735
       735
       +
                                   eprintln!("failed to join backup task: {e:?}");

     

       736
       736
       +
                               }

     

       737
       737
       +
                           }

     

       738
       738
       +
                       }

     

       739
       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
        
               }