add database backup with admin commands

Signed-off-by: Jason Volk <jason@zemos.net>

Cargo.lock

@@ -355,6 +355,15 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e"
 
+[[package]]
+name = "chrono"
+version = "0.4.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8eaf5903dcbc0a39312feb77df2ff4c76387d591b9fc7b04a238dcf8bb62639a"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "clang-sys"
 version = "1.7.0"
@@ -421,6 +430,7 @@ dependencies = [
  "axum-server-dual-protocol",
  "base64 0.22.0",
  "bytes",
+ "chrono",
  "clap",
  "cyborgtime",
  "either",

Cargo.toml

@@ -67,6 +67,12 @@ cyborgtime = "2.1.1"
 bytes = "1.5.0"
 http = "0.2.12"
 
+# standard date and time tools
+[dependencies.chrono]
+version = "0.4.35"
+features = ["alloc"]
+default-features = false
+
 # Web framework
 [dependencies.axum]
 version = "0.6.20"

src/config/mod.rs

@@ -42,6 +42,9 @@ pub struct Config {
 	#[serde(default = "default_database_backend")]
 	pub database_backend: String,
 	pub database_path: String,
+	pub database_backup_path: Option<String>,
+	#[serde(default = "default_database_backups_to_keep")]
+	pub database_backups_to_keep: i16,
 	#[serde(default = "default_db_cache_capacity_mb")]
 	pub db_cache_capacity_mb: f64,
 	#[serde(default = "default_new_user_displayname_suffix")]
@@ -252,6 +255,14 @@ impl fmt::Display for Config {
 			("Server name", self.server_name.host()),
 			("Database backend", &self.database_backend),
 			("Database path", &self.database_path),
+			(
+				"Database backup path",
+				match self.database_backup_path.as_ref() {
+					Some(path) => path,
+					None => "",
+				},
+			),
+			("Database backups to keep", &self.database_backups_to_keep.to_string()),
 			("Database cache capacity (MB)", &self.db_cache_capacity_mb.to_string()),
 			("Cache capacity modifier", &self.conduit_cache_capacity_modifier.to_string()),
 			("PDU cache capacity", &self.pdu_cache_capacity.to_string()),
@@ -446,6 +457,8 @@ fn default_port() -> ListeningPort {
 
 fn default_unix_socket_perms() -> u32 { 660 }
 
+fn default_database_backups_to_keep() -> i16 { 1 }
+
 fn default_database_backend() -> String { "rocksdb".to_owned() }
 
 fn default_db_cache_capacity_mb() -> f64 { 256.0 }

src/database/abstraction.rs

@@ -1,4 +1,4 @@
-use std::{future::Future, pin::Pin, sync::Arc};
+use std::{error::Error, future::Future, pin::Pin, sync::Arc};
 
 use super::Config;
 use crate::Result;
@@ -26,6 +26,14 @@ pub(crate) trait KeyValueDatabaseEngine: Send + Sync {
 
 	#[allow(dead_code)]
 	fn clear_caches(&self) {}
+
+	fn backup(&self) -> Result<(), Box<dyn Error>> {
+		unimplemented!()
+	}
+
+	fn backup_list(&self) -> Result<String> {
+		Ok(String::new())
+	}
 }
 
 pub(crate) trait KvTree: Send + Sync {

src/database/abstraction/rocksdb.rs

@@ -3,12 +3,17 @@ use std::{
 	pin::Pin,
 	sync::{Arc, RwLock},
 };
+use chrono::{
+	DateTime,
+	Utc,
+};
 
 use rust_rocksdb::{
+	backup::{BackupEngine, BackupEngineOptions},
 	LogLevel::{Debug, Error, Fatal, Info, Warn},
 	WriteBatchWithTransaction,
 };
-use tracing::{debug, info};
+use tracing::{debug, error, info};
 
 use super::{super::Config, watchers::Watchers, KeyValueDatabaseEngine, KvTree};
 use crate::{utils, Result};
@@ -221,6 +226,68 @@ impl KeyValueDatabaseEngine for Arc<Engine> {
 		Ok(())
 	}
 
+	fn backup(&self) -> Result<(), Box<dyn std::error::Error>> {
+		let path = self.config.database_backup_path.as_ref();
+		if path.is_none() || path.is_some_and(String::is_empty) {
+			return Ok(());
+		}
+
+		let options = BackupEngineOptions::new(&path.unwrap())?;
+		let mut engine = BackupEngine::open(&options, &self.env)?;
+		let ret = if self.config.database_backups_to_keep > 0 {
+			match engine.create_new_backup_flush(&self.rocks, true) {
+				Err(e) => return Err(Box::new(e)),
+				Ok(_) => {
+					let _info = engine.get_backup_info();
+					let info = &_info.last().unwrap();
+					info!(
+						"Created database backup #{} using {} bytes in {} files",
+						info.backup_id,
+						info.size,
+						info.num_files,
+					);
+					Ok(())
+				},
+			}
+		} else {
+			Ok(())
+		};
+
+		if self.config.database_backups_to_keep >= 0 {
+			let keep = u32::try_from(self.config.database_backups_to_keep)?;
+			if let Err(e) =  engine.purge_old_backups(keep.try_into()?) {
+				error!("Failed to purge old backup: {:?}", e.to_string())
+			}
+		}
+
+		ret
+	}
+
+	fn backup_list(&self) -> Result<String> {
+		let path = self.config.database_backup_path.as_ref();
+		if path.is_none() || path.is_some_and(String::is_empty) {
+			return Ok("Configure database_backup_path to enable backups".to_owned());
+		}
+
+		let mut res = String::new();
+		let options = BackupEngineOptions::new(&path.unwrap())?;
+		let engine = BackupEngine::open(&options, &self.env)?;
+		for info in engine.get_backup_info() {
+			std::fmt::write(&mut res, format_args!(
+				"#{} {}: {} bytes, {} files\n",
+				info.backup_id,
+				DateTime::<Utc>::from_timestamp(info.timestamp, 0)
+					.unwrap()
+					.to_rfc2822(),
+				info.size,
+				info.num_files,
+			))
+			.unwrap();
+		}
+
+		Ok(res)
+	}
+
 	// TODO: figure out if this is needed for rocksdb
 	#[allow(dead_code)]
 	fn clear_caches(&self) {}

src/database/key_value/globals.rs

@@ -272,4 +272,12 @@ lasttimelinecount_cache: {lasttimelinecount_cache}\n"
 		self.global.insert(b"version", &new_version.to_be_bytes())?;
 		Ok(())
 	}
+
+	fn backup(&self) -> Result<(), Box<dyn std::error::Error>> {
+		self.db.backup()
+	}
+
+	fn backup_list(&self) -> Result<String> {
+		self.db.backup_list()
+	}
 }

src/service/admin/mod.rs

@@ -436,6 +436,12 @@ enum ServerCommand {
 	ClearServiceCaches {
 		amount: u32,
 	},
+
+	/// - Backup the database
+	BackupDatabase,
+
+	/// - List database backups
+	ListBackups,
 }
 
 #[derive(Debug)]
@@ -1866,6 +1872,27 @@ impl Service {
 
 					RoomMessageEventContent::text_plain("Done.")
 				},
+				ServerCommand::ListBackups => {
+					let result = services().globals.db.backup_list()?;
+
+					RoomMessageEventContent::text_plain(result)
+				},
+				ServerCommand::BackupDatabase => {
+					let mut result = tokio::task::spawn_blocking(move || {
+						match services().globals.db.backup() {
+							Ok(_) => String::new(),
+							Err(e) => (*e).to_string(),
+						}
+					})
+					.await
+					.unwrap();
+
+					if result.is_empty() {
+						result = services().globals.db.backup_list()?;
+					}
+
+					RoomMessageEventContent::text_plain(&result)
+				},
 			},
 			AdminCommand::Debug(command) => match command {
 				DebugCommand::GetAuthChain {

src/service/globals/data.rs

@@ -1,4 +1,5 @@
 use std::collections::BTreeMap;
+use std::error::Error;
 
 use async_trait::async_trait;
 use ruma::{
@@ -32,4 +33,6 @@ pub trait Data: Send + Sync {
 	fn signing_keys_for(&self, origin: &ServerName) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>>;
 	fn database_version(&self) -> Result<u64>;
 	fn bump_database_version(&self, new_version: u64) -> Result<()>;
+	fn backup(&self) -> Result<(), Box<dyn Error>> { unimplemented!() }
+	fn backup_list(&self) -> Result<String> { Ok(String::new()) }
 }