// Copyright 2023 PingCAP, Inc. use std::{ borrow::Cow, cmp, collections::HashMap, fmt, fs, io::Write, path::PathBuf, sync::{ Arc, Mutex, RwLock, atomic::{AtomicU64, Ordering}, }, thread, time::Duration, }; use chrono::{DateTime, NaiveDateTime, TimeZone, Utc}; use cloud_encryption::MasterKey; use http::{Method, StatusCode, request::Parts}; use hyper::{Body, Response}; use kvengine::dfs::Dfs; use native_br::{ backup, backup::IncrementalBackupFile, backup_worker, backup_worker::BackupWorker, common::get_all_incremental_backups, online_snapshot::restore_keyspace_from_online_snapshot_with_cfg, restore, restore_keyspace, restore_keyspace::{ ObjectCache, ReportRestoreStepTrait, RestoreStep, RestoredKeyspace, restore_keyspace_with_cfg, }, }; use pd_client::{PdClient, pd_control::PdControl}; use serde::Deserialize; use tikv::storage::mvcc::TimeStamp; use tikv_util::{ AsyncSafeRuntime, HandyRwLock, config::{AbsoluteOrPercentSize, ReadableDuration, ReadableSize}, debug, error, errors::Context as _, info, sys::SysQuota, time::{Instant, Limiter}, warn, }; use crate::{ Config, common::{get_param, make_json_response, make_response}, error::{Error, Result}, metrics::{NATIVE_BR_COUNTER_VEC, NATIVE_BR_HISTOGRAM_VEC}, }; const MIN_PITR_INTERVAL_GAP_SECONDS: i64 = 1; // 1s const MAX_BACKUP_COUNT_PER_PAGE: usize = 1000; // Same with dfs list. const JSON_TIME_FORMAT: &str = "%Y-%m-%d %H:%M:%S%.3f"; // e.g. 2006-01-02 15:04:05.000 const BACKUP_NAME_FORMAT: &str = "%Y%m%d%H%M%S"; pub(crate) const BACKUPS_API_PATH: &str = "/api/v1/backups"; pub(crate) const RESTORE_KEYSPACE_API_PATH: &str = "/api/v1/restore_keyspace/"; const RESTORE_TASK_WORKING_PATH_PREFIX: &str = "r"; const ONLINE_SNAPSHOT_ATTEMPT_PATH: &str = "online-snapshot-attempt"; /// Backup and restore keyspace API: /// /// 1. backup: /// * GET /api/v1/backups?cluster_id=%d&last_backup_time= /// /// 2. restore keyspace: Specify backup_id&backup_name for normal restore or /// point_in_time for pitr. Online snapshot branch must be selected /// explicitly with restore_mode=online_snapshot and does not accept backup /// parameters. /// * PUT /api/v1/restore_keyspace/?cluster_id=%d&keyspace=%s& /// backup_id=%d&backup_name=%s[&source_keyspace=%s]& /// point_in_time= /// * GET /api/v1/restore_keyspace/?cluster_id=%d&keyspace=%s /// * DELETE /api/v1/restore_keyspace/?cluster_id=%d&keyspace=%s /// * GET /api/v1/restore_keyspace/?cluster_id=%d pub(crate) async fn handle_backup( manager: Arc, req: hyper::Request, ) -> hyper::Result> { let query = req.uri().query().unwrap_or(""); let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect(); match get_param::(&query_pairs, "cluster_id") { Some(cluster_id) if cluster_id == manager.get_cluster_id().unwrap() => {} _ => { return Ok(make_response( StatusCode::BAD_REQUEST, "Cluster ID mismatch", )); } } match *req.method() { Method::GET => { let ob_start_time = Instant::now(); let last_backup_time = query_pairs.get("last_backup_time").and_then(|s| { NaiveDateTime::parse_from_str(s, JSON_TIME_FORMAT) .ok() .map(|t| DateTime::::from_naive_utc_and_offset(t, Utc)) }); if last_backup_time.is_none() { return Ok(make_response( StatusCode::BAD_REQUEST, "Last backup time is none or invalid", )); } let start_backup_time = last_backup_time .unwrap() .checked_add_signed(chrono::Duration::seconds(1)) .unwrap(); let max_count = if let Some(count) = query_pairs.get("max_count") { match count.parse::() { Ok(c) => c, Err(e) => { return Ok(make_response( StatusCode::BAD_REQUEST, format!("max count is invalid {:?}", e), )); } } } else { MAX_BACKUP_COUNT_PER_PAGE }; match manager.list_backups(&start_backup_time, max_count).await { Ok((backups, has_more)) => { let resp = ListBackupResponse { items: backups.into_iter().map(Into::into).collect(), has_more, }; NATIVE_BR_HISTOGRAM_VEC .with_label_values(&["list_backup"]) .observe(ob_start_time.saturating_elapsed_secs()); Ok(make_json_response(StatusCode::OK, &resp)) } Err(e) => { NATIVE_BR_COUNTER_VEC .with_label_values(&["list_backup_fail"]) .inc(); Ok(make_response( StatusCode::NOT_FOUND, format!("Backups not found: {:?}", e), )) } } } _ => Ok(make_response(StatusCode::BAD_REQUEST, "Invalid method")), } } async fn parse_restore_source( manager: &NativeBrManager, query_pairs: &HashMap, Cow<'_, str>>, ) -> Result { let existed_backup = query_pairs.get("backup_id").is_some() || query_pairs.get("backup_name").is_some(); let pitr = query_pairs.get("point_in_time").is_some(); let restore_mode = query_pairs.get("restore_mode").map(|mode| mode.as_ref()); if let Some(mode) = restore_mode { if mode != "online_snapshot" { return Err(Error::CheckError(format!("Unknown restore mode: {mode}"))); } if existed_backup || pitr { return Err(Error::CheckError( "Online snapshot restore does not accept backup or point_in_time parameters" .to_string(), )); } return Ok(RestoreSource::OnlineSnapshot); } if existed_backup && pitr { return Err(Error::CheckError( "Request for normal restore and PiTR at the same time".to_string(), )); } if existed_backup { let backup_id = match get_param::(query_pairs, "backup_id") { Some(id) => id, None => { return Err(Error::CheckError("Backup ID is invalid".to_string())); } }; let backup_name = query_pairs .get("backup_name") .map(|s| s.to_string()) .unwrap_or_default(); let backup = IncrementalBackupFile::from_id(backup_id); if backup_name != backup.created_at().format(BACKUP_NAME_FORMAT).to_string() { return Err(Error::CheckError("Backup ID & name mismatch".to_string())); } return Ok(RestoreSource::Normal(backup)); } let ts = query_pairs .get("point_in_time") .map(|s| s.to_string()) .unwrap_or_default(); if ts.is_empty() { return Err(Error::CheckError("Recover time is empty".to_string())); } let utc_time = NaiveDateTime::parse_from_str(&ts, JSON_TIME_FORMAT) .map(|t| DateTime::::from_naive_utc_and_offset(t, Utc))?; // To avoid the system time gap between tikv-api and pd nodes. if Utc::now().signed_duration_since(utc_time).num_seconds() < MIN_PITR_INTERVAL_GAP_SECONDS { return Err(Error::CheckError(format!( "Future time {:?} is not supported", ts ))); } Ok(RestoreSource::Pitr { backup: manager.get_next_backup_after_ts(&utc_time).await?, point_in_time: utc_time, }) } pub(crate) async fn handle_restore_keyspace( manager: Arc, parts: Parts, ) -> hyper::Result> { let query = parts.uri.query().unwrap_or(""); let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect(); match get_param::(&query_pairs, "cluster_id") { Some(cluster_id) if cluster_id == manager.get_cluster_id().unwrap() => {} _ => { return Ok(make_response( StatusCode::BAD_REQUEST, "Cluster ID mismatch", )); } } let sub_path = parts .uri .path() .strip_prefix(RESTORE_KEYSPACE_API_PATH) .unwrap(); if sub_path.is_empty() && parts.method == Method::GET { return handle_get_all_restore_task(&manager); } let restore_id = match sub_path.parse::() { Ok(id) => id, Err(_) => { return Ok(make_response(StatusCode::BAD_REQUEST, "Invalid restore id")); } }; let target_keyspace = match query_pairs.get("keyspace") { Some(ks) if !ks.is_empty() => ks.to_string(), Some(_) => { return Ok(make_response( StatusCode::BAD_REQUEST, "Keyspace name is empty", )); } None => { return Ok(make_response( StatusCode::BAD_REQUEST, "Keyspace name is none", )); } }; // It's inplace restore for "keyspace" when "source_keyspace" is not specified. let source_keyspace = match query_pairs.get("source_keyspace") { Some(ks) if !ks.is_empty() => ks.to_string(), Some(_) => { return Ok(make_response( StatusCode::BAD_REQUEST, "Source keyspace name is empty", )); } None => target_keyspace.clone(), }; let keyspace_tag = format!("{}->{}", source_keyspace, target_keyspace); match parts.method { Method::GET => { debug!( "{} request to GET restore_keyspace, restore_id {}", keyspace_tag, restore_id ); handle_restore_status(&manager, restore_id, &target_keyspace) } Method::PUT => { let restore_source = match parse_restore_source(&manager, &query_pairs).await { Ok(restore_source) => restore_source, Err(err) => { return Ok(make_response( StatusCode::BAD_REQUEST, format!("Fail to get restore source: {}", err), )); } }; let restore_type = restore_source.restore_type(); debug!( "{} request to PUT restore_keyspace, restore_id {}, restore source {:?}", keyspace_tag, restore_id, restore_source ); let max_throughput = match query_pairs.get("max_throughput") { Some(v) => match v.parse::() { Ok(size) => size, Err(_) => { return Ok(make_response( StatusCode::BAD_REQUEST, format!("Invalid max_throughput value: {}", v), )); } }, None => ReadableSize(0), }; match manager.restore_keyspace( restore_id, source_keyspace, target_keyspace.clone(), restore_source.clone(), max_throughput, ) { Ok(true) => { info!( "{} restore keyspace started, restore_id {}, restore source {:?}", keyspace_tag, restore_id, restore_source ); let (progress, step_name) = RestoreProgressReporter::step_to_progress(&RestoreStep::Pending); let resp = RestoreProgressResponse { status: RestoreState::Pending, error: String::new(), duration: 0, id: restore_id, keyspace: target_keyspace, restore_type, restore_bytes: 0, tolerated_err: 0, restore_ts: 0, progress: RestoreProgress { step: step_name.to_string(), progress, }, }; Ok(make_json_response(StatusCode::OK, &resp)) } Ok(false) => { info!( "{} restore keyspace ignored, restore_id {}, restore source {:?}", keyspace_tag, restore_id, restore_source ); // Return current status when `restore_keyspace` request is ignored. handle_restore_status(&manager, restore_id, &target_keyspace) } Err(Error::RestoreKeyspaceTaskConflict(conflict_restore_id)) => { info!( "{} restore keyspace conflict, restore_id {}, restore source {:?}, conflict restore_id {}", keyspace_tag, restore_id, restore_source, conflict_restore_id, ); let resp = RestoreConflictResponse { keyspace: target_keyspace, id: restore_id, conflict_restore_id, }; Ok(make_json_response(StatusCode::CONFLICT, &resp)) } Err(e) => { error!( "{} restore keyspace error, restore_id {}, restore source {:?}, error {:?}", keyspace_tag, restore_id, restore_source, e ); handle_error(e) } } } Method::DELETE => { debug!( "{} request to DELETE restore_keyspace, restore_id {}", keyspace_tag, restore_id ); match manager.delete_restore(restore_id, &target_keyspace) { Ok(Some(true)) => { info!( "{} restore keyspace task deleted, restore_id {}", keyspace_tag, restore_id ); Ok(make_response(StatusCode::OK, "Restore task deleted")) } Ok(Some(false)) => { info!( "{} delete restore keyspace task ignored, restore_id {}", keyspace_tag, restore_id ); Ok(make_response( StatusCode::CONFLICT, "Restore task is not in final state", )) } Ok(None) => { info!( "{} delete restore keyspace task not found, restore_id {}", keyspace_tag, restore_id ); Ok(make_response( StatusCode::NOT_FOUND, format!("Restore task not found: {}", restore_id), )) } Err(err) => { error!( "{} delete restore keyspace task error, restore_id {}, error {:?}", keyspace_tag, restore_id, err ); handle_error(err) } } } _ => Ok(make_response(StatusCode::BAD_REQUEST, "Invalid method")), } } fn handle_get_all_restore_task(manager: &Arc) -> hyper::Result> { let resp_vec: Vec = manager .get_all_restore_task() .into_iter() .map(|(id, task)| RestoreProgressResponse::from_restore_status(id, task)) .collect(); Ok(make_json_response(StatusCode::OK, &resp_vec)) } fn handle_restore_status( manager: &Arc, restore_id: u64, keyspace: &str, ) -> hyper::Result> { match manager.restore_status(restore_id, keyspace) { Ok(Some(task)) => { let resp = RestoreProgressResponse::from_restore_status(restore_id, task); Ok(make_json_response(StatusCode::OK, &resp)) } Ok(None) => Ok(make_response( StatusCode::NOT_FOUND, format!("Restore task not found: {}", restore_id), )), Err(err) => { error!( "{} query restore keyspace task status error, restore_id {}, error {:?}", keyspace, restore_id, err ); handle_error(err) } } } fn handle_error(err: Error) -> hyper::Result> { match err { Error::CheckError(msg) => Ok(make_response(StatusCode::BAD_REQUEST, msg)), err => Ok(make_response( StatusCode::INTERNAL_SERVER_ERROR, err.to_string(), )), } } #[derive(Clone, Debug)] enum RestoreSource { Normal(IncrementalBackupFile), Pitr { backup: Option, point_in_time: DateTime, }, OnlineSnapshot, } impl RestoreSource { fn restore_type(&self) -> RestoreType { match self { Self::Normal(_) => RestoreType::Normal, Self::Pitr { .. } => RestoreType::Pitr, Self::OnlineSnapshot => RestoreType::OnlineSnapshot, } } } #[derive(Default, Serialize, Deserialize, Debug)] #[serde(default)] pub struct BackupItem { pub id: u64, pub name: String, pub time: String, } impl From for BackupItem { fn from(backup: IncrementalBackupFile) -> Self { Self { id: backup.id(), name: backup.created_at().format(BACKUP_NAME_FORMAT).to_string(), time: backup.created_at().format(JSON_TIME_FORMAT).to_string(), } } } #[derive(Default, Serialize, Deserialize, Debug)] #[serde(default)] pub struct ListBackupResponse { pub items: Vec, pub has_more: bool, } #[derive(Default, Serialize, Deserialize, Debug)] #[serde(default)] pub struct RestoreProgress { pub step: String, pub progress: i32, } #[derive(Default, Serialize, Deserialize, Debug)] #[serde(default)] pub struct RestoreProgressResponse { pub status: RestoreState, pub error: String, pub duration: i64, // in seconds pub id: u64, pub keyspace: String, pub restore_type: RestoreType, pub restore_bytes: u64, /// Number of unavailable stores tolerated by the completed restore. pub tolerated_err: usize, pub restore_ts: u64, pub progress: RestoreProgress, } impl RestoreProgressResponse { fn from_restore_status(restore_id: u64, status: RestoreTask) -> Self { let duration = status .end .unwrap_or_else(|| Utc::now()) .signed_duration_since(status.start); Self { status: status.state, error: status.error, duration: duration.num_seconds().clamp(0, i64::MAX), id: restore_id, keyspace: status.keyspace_name, restore_type: status.restore_type, restore_bytes: status.restore_bytes, tolerated_err: status.tolerated_err, restore_ts: status.restore_ts, progress: status.progress_reporter.get_progress(), } } } #[derive(Default, Serialize, Deserialize, Debug)] #[serde(default)] struct RestoreConflictResponse { keyspace: String, id: u64, conflict_restore_id: u64, } #[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq, PartialOrd)] pub enum RestoreState { #[default] Pending, Init, Running, Succeed, Error, } impl RestoreState { fn is_final(&self) -> bool { *self == Self::Succeed || *self == Self::Error } } #[derive(Clone, Copy, Default, Serialize, Deserialize, Debug, PartialEq)] pub enum RestoreType { #[default] Normal, Pitr, OnlineSnapshot, } #[derive(Clone)] pub(crate) struct RestoreTask { state: RestoreState, keyspace_name: String, error: String, start: DateTime, end: Option>, restore_type: RestoreType, // deprecated, remove after next upgrade. restore_params: RestoreParams, restore_bytes: u64, tolerated_err: usize, restore_ts: u64, progress_reporter: Arc, } impl RestoreTask { fn ttl_expired(&self, ttl: Duration) -> bool { self.state.is_final() && Utc::now() .signed_duration_since(self.end.unwrap()) .num_seconds() >= ttl.as_secs() as i64 } fn get_meta(&self) -> RestoreTaskMeta { RestoreTaskMeta { keyspace_name: self.keyspace_name.clone(), restore_type: self.restore_type, restore_params: self.restore_params.clone(), start: self.start.timestamp(), } } fn from_meta(meta: RestoreTaskMeta) -> Self { let start = meta.start(); Self { state: RestoreState::Error, // Consider as being interrupted by node restart. keyspace_name: meta.keyspace_name, error: "interrupted".to_string(), start, end: Some(Utc::now()), restore_type: meta.restore_type, restore_params: meta.restore_params.clone(), restore_bytes: 0, tolerated_err: 0, restore_ts: 0, progress_reporter: Arc::new(RestoreProgressReporter::new(0, RestoreStep::Init)), } } } #[derive(Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(default)] pub(crate) struct RestoreTaskMeta { keyspace_name: String, restore_type: RestoreType, // deprecated, remove after next upgrade. restore_params: RestoreParams, start: i64, // DateTime::timestamp(). } impl fmt::Debug for RestoreTaskMeta { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("RestoreTaskMeta") .field("keyspace_name", &self.keyspace_name) .field("restore_type", &self.restore_type) .field("restore_params", &self.restore_params) .field("start", &self.start()) .finish() } } impl RestoreTaskMeta { fn start(&self) -> DateTime { Utc.timestamp_opt(self.start, 0).unwrap() } } #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] #[serde(default)] struct RestoreParams { /// The backup name. For "Normal" restore. #[serde(skip_serializing_if = "Option::is_none")] backup_name: Option, /// The point in time. For "PiTR". #[serde(skip_serializing_if = "Option::is_none")] point_in_time: Option, /// The name of source keyspace for restore. /// /// For "Branch" or restore to another cluster. #[serde(skip_serializing_if = "Option::is_none")] source_keyspace: Option, } impl RestoreParams { fn new( restore_source: &RestoreSource, source_keyspace: &str, target_keyspace: &str, ) -> Result { let src_ks_opt = (source_keyspace != target_keyspace).then(|| source_keyspace.to_string()); match restore_source { RestoreSource::Normal(backup) => Ok(Self { backup_name: Some(backup.name().to_string()), source_keyspace: src_ks_opt, ..Default::default() }), RestoreSource::Pitr { point_in_time, .. } => Ok(Self { point_in_time: Some(point_in_time.timestamp()), source_keyspace: src_ks_opt, ..Default::default() }), RestoreSource::OnlineSnapshot => { let Some(source_keyspace) = src_ks_opt else { return Err(Error::CheckError( "Online snapshot restore requires different source and target keyspaces" .to_string(), )); }; Ok(Self { source_keyspace: Some(source_keyspace), ..Default::default() }) } } } #[inline] fn is_none(&self) -> bool { self.backup_name.is_none() && self.point_in_time.is_none() && self.source_keyspace.is_none() } } struct RestoreProgressReporter { restore_id: u64, step: Mutex, } impl RestoreProgressReporter { fn new(restore_id: u64, step: RestoreStep) -> Self { Self { restore_id, step: Mutex::new(step), } } fn step_to_progress(step: &RestoreStep) -> (i32, &'static str) { match step { // The difference value between this and next step is the estimated time cost of current // step. // Note: the message will be seen by customers. RestoreStep::Pending => (0, "pending"), RestoreStep::Init => (5, "init"), RestoreStep::InstantBackup => (10, "prepare backup"), RestoreStep::CaptureOnlineSnapshot => (10, "capture online snapshot"), RestoreStep::RemoveTiFlashReplicas => (30, "remove tiflash replicas"), RestoreStep::LoadBackupMeta => (35, "load backup meta"), RestoreStep::ExtractBackupShards => (40, "extract backup shards"), RestoreStep::ResolveLocks => (48, "resolve locks"), RestoreStep::FlushShards => (50, "flush shards"), RestoreStep::TruncateTs => (52, "truncate ts"), RestoreStep::SplitRegions => (55, "split regions"), RestoreStep::AlignRegions => (60, "align regions"), RestoreStep::RestoreSnapshotsToServers => (65, "restore snapshots to servers"), RestoreStep::RestoreSnapshotsFinished => (85, "restore snapshots finished"), RestoreStep::RetainSstFiles => (85, "retain files"), RestoreStep::Finalize => (90, "finalize"), // Wait for ClusterCR become normal. } } fn get_progress(&self) -> RestoreProgress { let step = self.step.lock().unwrap(); let (progress, step_name) = Self::step_to_progress(&step); RestoreProgress { step: step_name.to_string(), progress, } } } impl ReportRestoreStepTrait for RestoreProgressReporter { fn report_step(&self, step: RestoreStep) { info!("restore {} report_step: {:?}", self.restore_id, step); let mut current_step = self.step.lock().unwrap(); if *current_step < step { *current_step = step; } } } type TasksMap = HashMap; type KeyspacesMap = HashMap; pub(crate) struct BrContext { pub pd_client: Arc, pub data_dir: PathBuf, pub dfs: Arc, pub runtime: Arc, pub master_key: MasterKey, pub restore_tasks: RwLock, pub keyspace_tasks: RwLock, pub backup_worker: BackupWorker, restore_concurrency: usize, object_cache: Option, /// Global/shared rate limiter token bucket for native_br operations. /// /// Selection rules (mutually exclusive, NOT hierarchical): /// - If per-request `max_throughput > 0`, the request uses an independent /// token bucket (does not consume this global bucket). /// - If per-request `max_throughput == 0`, the request follows this global /// shared bucket (if enabled). /// /// Speed is set to `f64::INFINITY` when disabled (zero overhead). global_limiter: Limiter, pub v1x_tasks: RwLock>, } impl BrContext { async fn should_run_restore_instant_backup(&self) -> Result { if v1x::all_dfs_workers_disabled(&self.pd_client).await? { info!("restore keyspace: skip instant backup because DFS workers are disabled"); return Ok(false); } Ok(true) } fn make_byte_limiter( &self, max_throughput: ReadableSize, source: &'static str, request_id: u64, ) -> Option { if max_throughput.0 > 0 { info!( "native_br: create per-request {} limiter: {:.1} MiB/s, request_id: {}", source, max_throughput.as_mb_f64(), request_id, ); Some(native_br::limiter::ByteLimiter::new( max_throughput.0, source, request_id, )) } else { let speed = self.global_limiter.speed_limit(); if speed.is_finite() { info!( "native_br: create shared {} limiter: {:.1} MiB/s, request_id: {}", source, speed / (1024.0 * 1024.0), request_id, ); Some(native_br::limiter::ByteLimiter::from_limiter( self.global_limiter.clone(), source, request_id, )) } else { None } } } fn make_throughput_limiter( &self, rate_cfg: &native_br::limiter::RateLimitConfig, max_throughput: ReadableSize, request_id: u64, ) -> Option> { let byte_limiter = self.make_byte_limiter(max_throughput, "restore", request_id)?; match native_br::limiter::ThroughputLimiter::with_byte_limiter( byte_limiter, rate_cfg, self.pd_client.clone(), self.runtime.handle().clone(), ) { Ok(limiter) => Some(Arc::new(limiter)), Err(e) => { error!("native_br: failed to create throughput limiter: {:?}", e); None } } } fn update_global_limiter(&self, rate_cfg: &native_br::limiter::RateLimitConfig) { let new_speed = if rate_cfg.enable && rate_cfg.max_throughput.0 > 0 { rate_cfg.max_throughput.0 as f64 } else { if rate_cfg.enable && rate_cfg.max_throughput.0 == 0 { warn!( "native_br: restore_rate_limit.enable=true but max_throughput=0, treating as disabled" ); } f64::INFINITY }; let old_speed = self.global_limiter.speed_limit(); if new_speed != old_speed { info!( "native_br: update global rate limiter: {:.1} -> {:.1} MiB/s", old_speed / (1024.0 * 1024.0), new_speed / (1024.0 * 1024.0), ); self.global_limiter.set_speed_limit(new_speed); } } // Return false if the state is falling back. fn init_restore_state( &self, restore_id: u64, keyspace_name: &str, restore_type: RestoreType, restore_params: RestoreParams, ) -> Result { let new_task = |tasks: &mut TasksMap, restore_params: RestoreParams| -> Result { let mut keyspaces = self.keyspace_tasks.write().unwrap(); if let Some(&restore_id) = keyspaces.get(keyspace_name) { return Err(Error::RestoreKeyspaceTaskConflict(restore_id)); } keyspaces.insert(keyspace_name.to_owned(), restore_id); let task = RestoreTask { state: RestoreState::Init, keyspace_name: keyspace_name.to_string(), error: String::new(), start: Utc::now(), end: None, restore_type, restore_params, restore_bytes: 0, tolerated_err: 0, restore_ts: 0, progress_reporter: Arc::new(RestoreProgressReporter::new( restore_id, RestoreStep::Init, )), }; let task_meta = task.get_meta(); tasks.insert(restore_id, task); Ok(task_meta) }; let mut tasks = self.restore_tasks.write().unwrap(); Ok(match tasks.get_mut(&restore_id) { None => { let meta = new_task(&mut tasks, restore_params)?; self.persist_task_meta(restore_id, &meta)?; true } Some(task) if task.state == RestoreState::Error => { // Retry (Error -> Init): check_task(task, keyspace_name, Some(&restore_params))?; let meta = new_task(&mut tasks, restore_params)?; self.persist_task_meta(restore_id, &meta)?; true } Some(_) => false, }) } // Return false if the state is falling back. fn change_restore_state( &self, restore_id: u64, keyspace_name: &str, new_state: RestoreState, err: Option, restore_bytes: u64, restore_ts: u64, tolerated_err: usize, ) -> Result { let mut tasks = self.restore_tasks.write().unwrap(); match tasks.get_mut(&restore_id) { None => { debug_assert!(false); return Ok(false); } Some(task) => { check_task(task, keyspace_name, None)?; if task.state >= new_state { return Ok(false); } task.restore_bytes = restore_bytes; task.restore_ts = restore_ts; task.tolerated_err = tolerated_err; task.state = new_state; task.error = err.map(|err| format!("{:?}", err)).unwrap_or_default(); if task.state.is_final() { task.end = Some(Utc::now()); self.keyspace_tasks.write().unwrap().remove(keyspace_name); } } } Ok(true) } fn get_progress_reporter(&self, restore_id: u64) -> Option> { let tasks = self.restore_tasks.read().unwrap(); tasks .get(&restore_id) .map(|task| task.progress_reporter.clone()) } fn restore_keyspace_core( &self, config: Config, working_path: PathBuf, keyspace_name: &str, target_keyspace_name: &str, restore_source: RestoreSource, progress_reporter: Arc, max_throughput: ReadableSize, restore_id: u64, ) -> Result { let limiter = self.make_throughput_limiter( &config.native_br.restore_rate_limit, max_throughput, restore_id, ); if matches!(restore_source, RestoreSource::OnlineSnapshot) { return Ok(restore_keyspace_from_online_snapshot_with_cfg( config.to_restore_config(), keyspace_name, target_keyspace_name, working_path, self.dfs.clone(), self.pd_client.clone(), &self.runtime, Some(self.master_key.clone()), progress_reporter, limiter, )?); } // Preserve the current data for PiTR while DFS workers maintain a continuous // recovery chain. When every worker is disabled, that guarantee is already // unavailable, so an instant backup would provide no additional protection. let should_run_instant_backup = self .runtime .block_on(self.should_run_restore_instant_backup())?; let instant_backup = if should_run_instant_backup { progress_reporter.report_step(RestoreStep::InstantBackup); let instant_backup = self.runtime.block_on( self.backup_worker .instant_backup_with_retry(config.native_br.instant_backup_timeout.0), )?; info!("restore keyspace: instant backup: {:?}", instant_backup; "keyspace" => &keyspace_name, "target_keyspace" => &target_keyspace_name); Some(instant_backup) } else { None }; let (backup_file, truncate_ts) = match restore_source { RestoreSource::Normal(backup) => (backup, None), RestoreSource::Pitr { backup, point_in_time, } => { let backup = match backup { Some(backup) => backup, None => instant_backup .as_ref() .ok_or_else(|| { Error::CheckError( "cannot restore to the point in time without a later backup while DFS workers are disabled" .to_string(), ) })? .backup_file .clone(), }; ( backup, Some( TimeStamp::compose(point_in_time.timestamp_millis() as u64, 0).into_inner(), ), ) } RestoreSource::OnlineSnapshot => unreachable!("handled before instant backup"), }; let restore_config = config.to_restore_config(); Ok(restore_keyspace_with_cfg( restore_config, keyspace_name, target_keyspace_name, backup_file.name(), Some(working_path), self.dfs.clone(), self.pd_client.clone(), &self.runtime, truncate_ts, Some(self.master_key.clone()), progress_reporter, self.object_cache.clone(), limiter, )?) } fn restore_keyspace( &self, config: Config, restore_id: u64, keyspace_name: String, target_keyspace_name: String, restore_source: RestoreSource, max_throughput: ReadableSize, ) -> Result<()> { let ob_start_time = Instant::now(); self.change_restore_state( restore_id, &target_keyspace_name, RestoreState::Running, None, 0, 0, 0, )?; let progress_reporter = self.get_progress_reporter(restore_id).unwrap(); let working_path = if matches!(restore_source, RestoreSource::OnlineSnapshot) { let path = self .working_path(restore_id) .join(ONLINE_SNAPSHOT_ATTEMPT_PATH); match fs::remove_dir_all(&path) { Ok(()) => {} Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} Err(err) => { self.change_restore_state( restore_id, &target_keyspace_name, RestoreState::Error, Some(err.into()), 0, 0, 0, )?; return Ok(()); } } path } else { self.working_path(restore_id) }; let res = match self.restore_keyspace_core( config, working_path, &keyspace_name, &target_keyspace_name, restore_source, progress_reporter.clone(), max_throughput, restore_id, ) { Ok(ret) => { NATIVE_BR_COUNTER_VEC .with_label_values(&["restore_keyspace_succeed"]) .inc(); NATIVE_BR_HISTOGRAM_VEC .with_label_values(&["restore_keyspace"]) .observe(ob_start_time.saturating_elapsed_secs()); progress_reporter.report_step(RestoreStep::Finalize); let res = self.change_restore_state( restore_id, &target_keyspace_name, RestoreState::Succeed, None, ret.restore_bytes, ret.ts, ret.tolerated_err, ); self.remove_working_path(restore_id); res } Err(err) => { error!( "{}->{} restore_keyspace error, restore_id {}, error {:?}", keyspace_name, target_keyspace_name, restore_id, err ); NATIVE_BR_COUNTER_VEC .with_label_values(&["restore_keyspace_fail"]) .inc(); self.change_restore_state( restore_id, &target_keyspace_name, RestoreState::Error, Some(err), 0, 0, 0, ) } }; if let Err(e) = res { NATIVE_BR_COUNTER_VEC .with_label_values(&["change_restore_state_fail"]) .inc(); return Err(e); } Ok(()) } fn working_path(&self, restore_id: u64) -> PathBuf { self.data_dir .join(format!("{RESTORE_TASK_WORKING_PATH_PREFIX}{restore_id}")) } fn remove_working_path(&self, restore_id: u64) { let working_path = self.working_path(restore_id); if let Err(e) = fs::remove_dir_all(&working_path) { warn!("fail to remove working path: {:?}", e; "path" => working_path.display()); } } fn persist_task_meta(&self, restore_id: u64, meta: &RestoreTaskMeta) -> Result<()> { static TMP_ID: AtomicU64 = AtomicU64::new(0); let working_path = self.working_path(restore_id); fs::create_dir_all(&working_path).ctx("create_working_path")?; let tmp_path = working_path.join(format!( "meta.json.{}", TMP_ID.fetch_add(1, Ordering::Relaxed) )); let mut tmp = fs::File::create(&tmp_path).ctx("create_tmp")?; tmp.write_all(&serde_json::to_vec(&meta)?) .ctx("write_meta")?; tmp.sync_data().ctx("sync")?; drop(tmp); let meta_file = working_path.join("meta.json"); fs::rename(tmp_path, meta_file).ctx("rename")?; Ok(()) } fn read_task_meta(&self, restore_id: u64) -> Result { let working_path = self.working_path(restore_id); let meta_file = working_path.join("meta.json"); if !meta_file.exists() { return Err(Error::CheckError(format!( "Meta file not found: {}", meta_file.display() ))); } let meta = fs::read(&meta_file).ctx("read_meta")?; let meta = serde_json::from_slice::(&meta)?; Ok(meta) } fn init(&mut self) -> Result<()> { let entries = fs::read_dir(&self.data_dir).ctx("read_dir")?; for entry in entries { let entry = entry.ctx("entry")?; let path = entry.path(); if path.is_dir() && path.file_name().is_some() { let name = path.file_name().unwrap().to_string_lossy(); let name = name.as_ref(); if let Some(restore_id) = name .strip_prefix(RESTORE_TASK_WORKING_PATH_PREFIX) .and_then(|s| s.parse::().ok()) { match self.read_task_meta(restore_id) { Ok(meta) => { info!("restore task: {:?}", meta); let task = RestoreTask::from_meta(meta); self.restore_tasks.wl().insert(restore_id, task); } Err(err) => { warn!("fail to read task meta: {:?}", err; "restore_id" => restore_id); } } } } } Ok(()) } } impl Drop for BrContext { fn drop(&mut self) { self.backup_worker.stop(); } } #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(default)] #[serde(rename_all = "kebab-case")] pub struct NativeBrConfig { /// The time-to-live when restore task has been in final state. pub restore_task_ttl: ReadableDuration, /// The timeout for waiting the flush of mem-tables. pub restore_timeout_wait_flush: ReadableDuration, /// The timeout for the requests of restoring snapshots to TiKV servers. pub restore_timeout_restore_snapshot: ReadableDuration, pub restore_snapshot_concurrency: usize, /// The timeout for fetching the latest wal chunk from store. pub restore_timeout_fetch_wal: ReadableDuration, pub restore_timeout_pd_control: ReadableDuration, /// The maximum number of retries for the process from split regions to /// restore snapshots. pub restore_max_retry: usize, pub restore_coarse_split_regions_factor: usize, pub online_snapshot_max_unaligned_voter_peers: usize, pub online_snapshot_attempt_timeout: ReadableDuration, /// Maximum number of source Regions captured concurrently by an online /// snapshot branch attempt. pub online_snapshot_capture_concurrency: usize, /// Maximum number of target Regions aligned concurrently by an online /// snapshot branch attempt. pub online_snapshot_placement_concurrency: usize, /// The timeout for instant backup. pub instant_backup_timeout: ReadableDuration, /// Interval for periodic backup. `0s` to disable periodic backup. pub backup_interval: ReadableDuration, /// Interval for gathering backup requests as a batch. pub backup_batch_interval: ReadableDuration, /// Delay before perform backup. Also used as the switch for passing /// `backup_ts` to rfengine during backup. pub backup_delay: ReadableDuration, pub backup_ts_wait_timeout: ReadableDuration, pub backup_ts_ttl: ReadableDuration, pub backup_pd_meta_watch_cache_enabled: bool, pub backup_pd_meta_watch_cache_progress_timeout: ReadableDuration, #[cfg(feature = "testexport")] pub backup_skip_keyspace_meta: bool, /// Whether to tolerate unavailability of no more than one store when /// backup. pub backup_tolerate_err: bool, /// Whether to tolerate unavailability of no more than one store when /// restore. pub restore_tolerate_err: bool, /// Whether restore may tolerate only explicitly supported transient errors. pub restore_strict_tolerate: bool, pub restore_store_concurrency: usize, pub restore_concurrency_per_core: f64, pub restore_object_cache_capacity: AbsoluteOrPercentSize, pub restore_cache_for_decompressed_wal_chunks: bool, pub wal_chunk_target_file_size: ReadableSize, /// See `RestoreConfig::lower_memory`. pub lower_memory: bool, /// Default/global rate limit shared across all native_br operations /// (restore/copy/unpack). /// /// Selection rules (mutually exclusive, NOT hierarchical): /// - If per-request `max_throughput > 0`, the request uses an independent /// token bucket (does not consume this global bucket). /// - If per-request `max_throughput == 0`, the request follows this global /// shared bucket (if enabled). pub restore_rate_limit: native_br::limiter::RateLimitConfig, } impl Default for NativeBrConfig { fn default() -> Self { Self { restore_task_ttl: ReadableDuration::minutes(60), restore_timeout_wait_flush: restore::DEFAULT_TIMEOUT_WAIT_FLUSH, restore_timeout_restore_snapshot: restore::DEFAULT_TIMEOUT_RESTORE_SNAPSHOT, restore_snapshot_concurrency: restore::DEFAULT_RESTORE_SNAPSHOT_CONCURRENCY, restore_timeout_fetch_wal: restore::DEFAULT_TIMEOUT_FETCH_WAL, restore_timeout_pd_control: restore::DEFAULT_TIMEOUT_PD_CONTROL, restore_max_retry: restore::DEFAULT_RESTORE_MAX_RETRY, restore_coarse_split_regions_factor: 64, online_snapshot_max_unaligned_voter_peers: restore::DEFAULT_ONLINE_SNAPSHOT_MAX_UNALIGNED_VOTER_PEERS, online_snapshot_attempt_timeout: restore::DEFAULT_ONLINE_SNAPSHOT_ATTEMPT_TIMEOUT, online_snapshot_capture_concurrency: restore::DEFAULT_ONLINE_SNAPSHOT_CAPTURE_CONCURRENCY, online_snapshot_placement_concurrency: restore::DEFAULT_ONLINE_SNAPSHOT_PLACEMENT_CONCURRENCY, instant_backup_timeout: backup_worker::DEFAULT_TIMEOUT_INSTANT_BACKUP, backup_interval: ReadableDuration::ZERO, backup_batch_interval: ReadableDuration::secs(15), backup_delay: ReadableDuration::ZERO, backup_ts_wait_timeout: backup::BACKUP_TS_WAIT_TIMEOUT_DEFAULT, backup_ts_ttl: backup::BACKUP_TS_TTL_DEFAULT, backup_pd_meta_watch_cache_enabled: false, backup_pd_meta_watch_cache_progress_timeout: ReadableDuration::secs(2), #[cfg(feature = "testexport")] backup_skip_keyspace_meta: false, backup_tolerate_err: false, restore_tolerate_err: false, restore_strict_tolerate: false, restore_store_concurrency: restore_keyspace::RESTORE_RFENGINE_CONCURRENCY, restore_concurrency_per_core: 1.0, restore_object_cache_capacity: 0.into(), restore_cache_for_decompressed_wal_chunks: false, wal_chunk_target_file_size: ReadableSize::mb(64), lower_memory: false, restore_rate_limit: native_br::limiter::RateLimitConfig::default(), } } } impl NativeBrConfig { fn object_cache_average_item_size(&self) -> u64 { if self.restore_cache_for_decompressed_wal_chunks { self.wal_chunk_target_file_size.0 } else { self.wal_chunk_target_file_size.0 / 3 } } } pub(crate) struct NativeBrManager { pub context: Arc, pub config: RwLock, } impl NativeBrManager { pub(crate) fn new( runtime: Arc, pd_client: Arc, dfs: Arc, data_dir: PathBuf, master_key: MasterKey, config: Config, ) -> Self { let backup_config = config.to_backup_config(); let backup_worker = BackupWorker::new( backup_config, pd_client.clone(), config.native_br.backup_interval.0, config.native_br.backup_batch_interval.0, ); let restore_concurrency = cmp::max( (SysQuota::cpu_cores_quota() * config.native_br.restore_concurrency_per_core) as usize, 2, ); info!( "native_br: restore concurrency limit: {}", restore_concurrency ); let object_cache_capacity = config .native_br .restore_object_cache_capacity .as_memory_size(); let object_cache = (object_cache_capacity > 0).then(|| { info!("native_br: create object cache"; "capacity" => object_cache_capacity); ObjectCache::new( object_cache_capacity, config.native_br.object_cache_average_item_size(), ) }); let global_limiter = { let rate_cfg = &config.native_br.restore_rate_limit; let speed = if rate_cfg.enable && rate_cfg.max_throughput.0 > 0 { info!( "native_br: global rate limiter enabled: {:.1} MiB/s", rate_cfg.max_throughput.as_mb_f64() ); rate_cfg.max_throughput.0 as f64 } else { f64::INFINITY }; ::builder(speed) .refill(Duration::from_millis(10)) .build() }; let mut context = BrContext { pd_client, dfs, data_dir, runtime, master_key, restore_tasks: Default::default(), keyspace_tasks: Default::default(), backup_worker, restore_concurrency, object_cache, global_limiter, v1x_tasks: Default::default(), }; if let Err(err) = context.init() { warn!("BR context init failed: {:?}", err); } Self { context: Arc::new(context), config: RwLock::new(config), } } pub(crate) fn update_native_br_config(&self, config: NativeBrConfig) { let mut ori_config = self.config.write().unwrap(); if ori_config.native_br != config { info!( "Update native br config from {:?} to {:?}", ori_config.native_br, config ); self.context .update_global_limiter(&config.restore_rate_limit); ori_config.native_br = config; } } async fn list_backups( &self, start_backup_time: &DateTime, max_count: usize, ) -> Result<(Vec, bool)> { let (backups, has_more) = get_all_incremental_backups( self.context.dfs.as_ref(), &start_backup_time.date_naive(), Some(&start_backup_time.time()), max_count, ) .await?; Ok((backups, has_more)) } fn get_cluster_id(&self) -> Result { Ok(self.context.pd_client.get_cluster_id()?) } /// Return: /// Ok(true): task started. /// Ok(false): request ignored due to duplicated. /// Err(err): error occurred. fn restore_keyspace( &self, restore_id: u64, keyspace_name: String, target_keyspace_name: String, restore_source: RestoreSource, max_throughput: ReadableSize, ) -> Result { let current_tasks = self.get_not_final_task_count(); let limit = self.context.restore_concurrency; if current_tasks >= limit { warn!("reach concurrency limit"; "current" => current_tasks, "limit" => limit); // TODO: pending in a queue, other than return error directly. return Err(Error::ReachConcurrencyLimit { current: current_tasks, limit, }); } let restore_type = restore_source.restore_type(); let restore_params = RestoreParams::new(&restore_source, &keyspace_name, &target_keyspace_name)?; if self.context.init_restore_state( restore_id, &target_keyspace_name, restore_type, restore_params, )? { let context = self.context.clone(); let config = self.config.read().unwrap().clone(); thread::spawn(move || { context.restore_keyspace( config, restore_id, keyspace_name, target_keyspace_name, restore_source, max_throughput, ) }); Ok(true) } else { Ok(false) } } /// Return: /// Ok(Some): succeed. /// Ok(None): not found. /// Err(err): error occurred. fn restore_status(&self, restore_id: u64, keyspace_name: &str) -> Result> { let tasks = self.context.restore_tasks.rl(); if let Some(task) = tasks.get(&restore_id) { check_task(task, keyspace_name, None)?; Ok(Some(task.clone())) } else { Ok(None) } } /// Return all restore tasks in memory. fn get_all_restore_task(&self) -> HashMap { self.context.restore_tasks.rl().clone() } /// Return number of tasks which are in not final state. fn get_not_final_task_count(&self) -> usize { self.context .restore_tasks .rl() .values() .filter(|task| !task.state.is_final()) .count() } /// Return: /// Ok(Some(true)): deleted. /// Ok(Some(false)): ignored due state is not final. /// Ok(None): restore task not found. /// Err(err): error occurred. fn delete_restore(&self, restore_id: u64, keyspace_name: &str) -> Result> { let mut tasks = self.context.restore_tasks.wl(); if let Some(task) = tasks.get(&restore_id) { check_task(task, keyspace_name, None)?; if task.state.is_final() { tasks.remove(&restore_id); self.context.remove_working_path(restore_id); Ok(Some(true)) } else { Ok(Some(false)) } } else { Ok(None) } } pub fn cleanup_expired_restores(&self) { let ttl = self.config.rl().native_br.restore_task_ttl.0; let restores = { let tasks = self.context.restore_tasks.rl(); tasks .iter() .filter_map(|(&restore_id, task)| { task.ttl_expired(ttl) .then_some((restore_id, task.keyspace_name.clone())) }) .collect::>() }; for (restore_id, keyspace_name) in restores { if let Ok(Some(true)) = self.delete_restore(restore_id, &keyspace_name) { info!("{}({}) expired and removed", keyspace_name, restore_id); } } } async fn get_next_backup_after_ts( &self, ts: &DateTime, ) -> Result> { let (backup_files, _) = self.list_backups(ts, 1).await?; Ok(backup_files.first().cloned()) } #[allow(dead_code)] pub fn get_pd_ctl(&self) -> Result { let cfg = self.config.rl().pd.clone(); let sec_mgr = self.context.pd_client.get_security_mgr(); let pd_ctl = PdControl::new(cfg, sec_mgr)?; Ok(pd_ctl) } } fn check_task( task: &RestoreTask, keyspace_name: &str, request_params: Option<&RestoreParams>, ) -> Result<()> { if task.keyspace_name != keyspace_name { return Err(Error::CheckError( "restore id & keyspace not match".to_string(), )); } if let Some(request_params) = request_params { check_task_params(task, request_params)?; } Ok(()) } fn check_task_params(task: &RestoreTask, request_params: &RestoreParams) -> Result<()> { // `task.restore_params.is_none()` is true when startup from meta of old // version. if task.restore_params.is_none() || &task.restore_params == request_params { Ok(()) } else { let msg = format!( "restore params not match: task: {:?}, request: {:?}", task.restore_params, request_params ); Err(Error::CheckError(msg)) } } pub mod v1x { use std::{ collections::HashMap, fmt::{Debug, Display}, marker::PhantomData, ops::Add, path::Path, sync::Arc, time::Duration, }; use chrono::{DateTime, DurationRound, NaiveDate, NaiveDateTime, Utc}; use futures::{StreamExt, TryStreamExt, stream::FuturesUnordered}; use http::{Method, StatusCode}; use hyper::{Body, Request, Response}; use kvengine::dfs::{DFSConfig, Dfs}; use kvenginepb::PackedBackup; use kvproto::metapb::Store; use native_br::{ archive::get_latest_archive_date, backup::{IncrementalBackupFile, backup_file_full_path}, common::{ INCREMENTAL_BACKUP_FOLDER_FORMAT, get_all_incremental_backups, send_request_to_store, }, error::Error as NativeBrError, packing::{ CopyPackedRun, CopyStep, MigratePackEnv, PackBackupStep, PackConfig, PackContext, ReportCopyStepTrait, ReportPackBackupStepTrait, ReportUnpackStepTrait, UnpackRun, UnpackStep, offline_pd::OfflinePd, }, restore_keyspace::{self, ReportRestoreStepTrait, RestoreStep}, }; use pd_client::{PdClient, pd_control::PdControl}; use protobuf::Message; use rfenginepb::ClusterBackupMeta; use serde::{ Deserialize, Deserializer, Serialize, de::{ IntoDeserializer, value::{MapDeserializer, StrDeserializer}, }, }; use serde_json::json; use tikv::storage::mvcc::TimeStamp; use tikv_util::{HandyRwLock, box_err, box_try, config::ReadableSize, defer, info, warn}; use tokio::task::spawn_blocking; use crate::{ common::make_json_response, error::Error, metrics::{self, NATIVE_BR_HISTOGRAM_VEC, NATIVE_BR_V1X_COUNTER_VEC}, native_br::{BACKUP_NAME_FORMAT, BackupItem, BrContext, NativeBrManager}, }; pub(crate) const API_V1X: &str = "/api/v1/x/"; type HttpResult = std::result::Result; type Result = std::result::Result; fn ts_to_datetime(ts: impl Into) -> DateTime { let ts: TimeStamp = ts.into(); NaiveDateTime::from_timestamp_millis(ts.physical() as _) .unwrap_or_else(|| panic!("timestamp too huge to be converted to datetime: {}", ts)) .and_utc() } struct AsyncDropDfs(Option>); impl From> for AsyncDropDfs { fn from(dfs: Arc) -> Self { Self(Some(dfs)) } } impl std::ops::Deref for AsyncDropDfs { type Target = Arc; fn deref(&self) -> &Self::Target { self.0.as_ref().unwrap() } } impl Drop for AsyncDropDfs { fn drop(&mut self) { let ptr = self.0.take().unwrap(); if let Ok(h) = tokio::runtime::Handle::try_current() { h.spawn_blocking(move || drop(ptr)); } // Or... We are not in a tokio context. It is safe to drop the // runtime. } } #[derive(Debug)] pub struct HttpError(Response); fn internal_code(e: &Error) -> &'static str { match e { Error::CheckError(_) => "CheckErr", Error::NotFound { .. } => "NotFound", Error::RestoreKeyspaceTaskConflict(_) => "RestoreKeyspaceTaskConflict", Error::Existed(_) => "Existed", Error::InvalidStateTrans { .. } => "InvalidStateTrans", Error::ReachConcurrencyLimit { .. } => "ReachConcurrencyLimit", _ => "InternalErr", } } impl> From for HttpError { fn from(err: T) -> Self { let crate_err: Error = err.into(); let ic = internal_code(&crate_err); match crate_err { Error::CheckError(msg) => Self::error_response(ic, StatusCode::BAD_REQUEST, msg), Error::NotFound { resource, identity, notes, } => Self::error_response( ic, StatusCode::NOT_FOUND, json!({ "resource": resource, "identity": identity, "notes": notes }), ), Error::RestoreKeyspaceTaskConflict(task_id) => { Self::error_response(ic, StatusCode::CONFLICT, json!({"task_id": task_id})) } Error::Existed(msg) => Self::error_response( ic, StatusCode::CONFLICT, format!("resource already existed: {}", msg), ), Error::InvalidStateTrans { resource, from, to } => Self::error_response( ic, StatusCode::CONFLICT, json!({ "resource": resource, "from": from, "to": to, }), ), Error::ReachConcurrencyLimit { limit, .. } => Self::error_response( ic, StatusCode::TOO_MANY_REQUESTS, format!("concurrent running task reaches limit: {}", limit), ), otherwise => Self::error_response( ic, StatusCode::INTERNAL_SERVER_ERROR, otherwise.to_string(), ), } } } impl serde::de::Error for Error { fn custom(msg: T) -> Self where T: Display, { Error::CheckError(msg.to_string()) } } impl std::error::Error for HttpError {} impl Display for HttpError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, " (body = {:?})", self.0.status(), self.0.body() ) } } impl HttpError { fn error_response( internal_code: impl ToString, status_code: StatusCode, body: impl Serialize, ) -> Self { assert!(!status_code.is_success(), "{}", status_code); let b = json!({ "internal_code": internal_code.to_string(), "error": body, }); Self(make_json_response(status_code, &b)) } } struct FieldHinttedStrDe<'a, E> { content: &'a str, hint: &'a str, phantom: PhantomData, } impl<'a, E> FieldHinttedStrDe<'a, E> { fn new(content: &'a str, hint: &'a str) -> Self { Self { content, hint, phantom: PhantomData, } } } impl IntoDeserializer<'_, E> for FieldHinttedStrDe<'_, E> { type Deserializer = Self; fn into_deserializer(self) -> Self::Deserializer { self } } impl<'de, E: serde::de::Error> Deserializer<'de> for FieldHinttedStrDe<'_, E> { type Error = E; fn deserialize_any(self, visitor: V) -> std::result::Result where V: serde::de::Visitor<'de>, { visitor.visit_str(self.content).map_err(|err: Self::Error| { ::custom(format!( "during handing key {}: {}", self.hint, err )) }) } fn deserialize_option( self, visitor: V, ) -> std::prelude::v1::Result where V: serde::de::Visitor<'de>, { if self.content.is_empty() { visitor.visit_none() } else { visitor.visit_some(StrDeserializer::new(self.content)) } .map_err(|err: Self::Error| { ::custom(format!( "during handing key {}: {}", self.hint, err )) }) } serde::forward_to_deserialize_any! { bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes byte_buf unit unit_struct newtype_struct seq tuple tuple_struct map struct enum identifier ignored_any } } fn v1_compat_info(key: &str) -> Option<(u64, String)> { IncrementalBackupFile::try_from_full_path(key).map(|file| { ( file.id(), file.created_at().format(BACKUP_NAME_FORMAT).to_string(), ) }) } #[derive(Deserialize, Serialize)] pub struct ListBackupRequest { #[serde(default)] from_prefix: String, #[serde(with = "serde_with::rust::display_fromstr", default)] load_details: bool, #[serde( with = "serde_with::rust::display_fromstr", default = "default_max_count" )] max_count: u64, #[serde(flatten)] s3_override: S3Override, } fn default_max_count() -> u64 { 1 } #[derive(Serialize, Deserialize, Debug)] pub struct ListBackupResponse { pub data: Vec, pub cluster_id: u64, pub has_more: bool, } #[derive(Debug, Serialize, Deserialize)] pub struct Backup { pub name: String, pub last_modify_time: String, #[serde(skip_serializing_if = "Option::is_none")] pub details: Option, pub v1_compat_id: Option, pub v1_compat_name: Option, } impl Backup { pub fn try_to_backup_item(&self) -> Option { IncrementalBackupFile::try_from_full_path(&format!("/backup/{}", self.name)) .map(From::from) } } #[derive(Debug, Default, Serialize, Deserialize)] pub struct KeyspaceBackupInfo { pub keyspace_id: u32, // Skip this field when failed to fetch it. // So the caller knows there are something wrong... #[serde(skip_serializing_if = "String::is_empty")] pub keyspace_name: String, pub keyspace_state: String, pub size: u64, } #[derive(Debug, Serialize, Deserialize)] pub struct BackupDetails { pub backup_ts: u64, pub safe_ts: u64, pub keyspaces: Vec, /// Number of keyspaces skipped because their metadata could not be /// parsed. #[serde(default)] pub skipped_keyspace_count: usize, } #[derive(Debug, Default, Serialize, Deserialize)] pub struct FlushResult {} #[derive(Debug, Serialize, Deserialize)] pub struct CreateLightweightBackupResponse { pub backup_name: String, } #[derive(Serialize, Deserialize, Clone, Debug)] #[serde(tag = "state")] pub enum TaskState { Pending, Init, Running { info_last_update_time: DateTime, step: String, info: serde_json::Value, }, Error { error: String, }, Done { info: serde_json::Value, }, } #[derive(Serialize, Deserialize, Debug, Clone)] struct StateTransform { to: State, at: DateTime, } impl TaskState { pub fn is_done(&self) -> bool { matches!(self, TaskState::Done { .. }) } pub fn state_string(&self) -> String { let mut base = serde_json::to_value(self).unwrap()["state"] .as_str() .unwrap() .to_owned(); if let Self::Running { step: status, .. } = self { base.push_str("::"); base.push_str(status.as_str()); } base } fn init(&mut self) { if let TaskState::Done { .. } | TaskState::Running { .. } = self { warn!("invalid state transform: Done | Running -> Init"); return; } if let TaskState::Error { error } = self { info!("retrying failed task"; "err" => %error); } *self = TaskState::Init; } fn running(&mut self, to_step: String) { if let TaskState::Done { .. } | TaskState::Error { .. } = self { warn!("invalid state transform: Done | Error -> Running"); return; } *self = Self::Running { info_last_update_time: Utc::now(), step: to_step, info: serde_json::Value::Null, } } fn running_info(&mut self, to_step: String, info: impl Serialize) { if let TaskState::Done { .. } | TaskState::Error { .. } = self { warn!("invalid state transform: Done | Error -> Running"); return; } *self = Self::Running { info_last_update_time: Utc::now(), step: to_step, info: serde_json::to_value(info).unwrap_or_default(), }; } fn error(&mut self, err: &Error) { if let TaskState::Error { .. } | TaskState::Done { .. } = self { warn!("invalid state transform: Error | Done -> Error"; "err" => %err); return; } *self = Self::Error { error: err.to_string(), }; } fn done(&mut self, by: serde_json::value::Value) { if let TaskState::Error { .. } | TaskState::Done { .. } = self { warn!("invalid state transform: Error | Done -> Done"; "by" => %by); return; } *self = Self::Done { info: by }; } } #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] #[serde(tag = "type")] pub enum TaskRequest { PackBackup(CreatePackBackupRequest), RestorePacked(CreateRestorePackedBackupRequest), CopyPacked(CreateCopyBackupRequest), } impl TaskRequest { fn id(&self) -> u64 { match self { TaskRequest::PackBackup(r) => r.id, TaskRequest::RestorePacked(r) => r.id, TaskRequest::CopyPacked(r) => r.id, } } fn task_type(&self) -> &'static str { match self { TaskRequest::PackBackup(_) => "pack_backup", TaskRequest::RestorePacked(_) => "restore_packed", TaskRequest::CopyPacked(_) => "copy_packed", } } async fn run(self, br_mgr: Arc) { let m = metrics::NATIVE_BR_V1X_BACKGROUND_TASKS.with_label_values(&[self.task_type()]); m.inc(); defer! { m.dec() } match self { TaskRequest::PackBackup(req) => Self::run_create_pack_backup(br_mgr, req).await, TaskRequest::RestorePacked(req) => Self::run_restore_packed(br_mgr, req).await, TaskRequest::CopyPacked(req) => Self::run_copy_packed(br_mgr, req).await, } } async fn run_copy_packed(br: Arc, req: CreateCopyBackupRequest) { let br_cx = &br.context; let dfs = br.overridden_storage_full( req.exotic.exotic_backup_endpoint.as_deref(), &req.s3_override, ); let exec = async { let reporter = Arc::new(ApiV1xReporter { cx: br_cx.clone(), id: req.id, }) as Arc; reporter.report_step(CopyStep::FetchMeta); let mig_env = MigratePackEnv::load_exotic(Arc::clone(&*dfs), &req.exotic.exotic_backup) .await?; let mig_env = match br_cx.make_byte_limiter(req.max_throughput, "copy", req.id) { Some(bl) => mig_env.with_byte_limiter(Arc::new(bl)), None => mig_env, }; let mut copy_run = CopyPackedRun::new(mig_env, &req.copy_to, reporter.clone()); let res = copy_run.execute().await?; Result::Ok(res) }; br_cx.finish_task(req.id, exec.await) } async fn run_restore_packed( br: Arc, req: CreateRestorePackedBackupRequest, ) { let cfg = br.config.rl().clone(); let br_cx = &br.context; let dfs = br.overridden_storage(req.exotic.exotic_backup_endpoint.as_deref()); let pd_cli = br_cx.pd_client.clone(); let exec = async { let truncate_ts = req .point_in_time .map(|t| TimeStamp::compose(t.naive_utc().timestamp_millis() as _, 0)); let mig_env = MigratePackEnv::load_exotic(dfs.clone(), &req.exotic.exotic_backup).await?; let mig_env = match br_cx.make_byte_limiter(req.max_throughput, "unpack", req.id) { Some(bl) => mig_env.with_byte_limiter(Arc::new(bl)), None => mig_env, }; let packed_keyspace_id = mig_env.get_packed_backup().keyspace_id; let packed_keyspace_name = mig_env.get_packed_backup().keyspace_name.clone(); let packed_cluster_id = mig_env.get_packed_backup().cluster_id; if truncate_ts .is_some_and(|ts| ts.into_inner() < mig_env.get_packed_backup().safe_ts) { return Err(Error::CheckError(format!( "you want to restore to {} but it was already reaped by GC (until {})", ts_to_datetime(truncate_ts.unwrap()), ts_to_datetime(mig_env.get_packed_backup().safe_ts), ))); } let reporter = Arc::new(ApiV1xReporter { cx: br_cx.clone(), id: req.id, }); let limiter = br_cx.make_throughput_limiter( &cfg.native_br.restore_rate_limit, req.max_throughput, req.id, ); let mut run = UnpackRun::new(mig_env, pd_cli.clone(), reporter.clone()); let target = run.execute().await?; let mut restore_cfg = cfg.to_restore_config(); restore_cfg.restore_packed_backup = true; let pd_control = PdControl::new(restore_cfg.pd.clone(), pd_cli.get_security_mgr())?; let target_keyspace = pd_control.get_keyspace_by_name(&req.keyspace).await?; let working_path = br_cx.working_path(req.id); let br_cx = br_cx.clone(); let dfs = Arc::clone(&*dfs); if br_cx.should_run_restore_instant_backup().await? { ReportRestoreStepTrait::report_step( reporter.as_ref(), RestoreStep::InstantBackup, ); br_cx .backup_worker .instant_backup_with_retry(cfg.native_br.instant_backup_timeout.0) .await?; } let restored = spawn_blocking(move || { restore_keyspace::restore_keyspace( packed_keyspace_id, target_keyspace.id, &target, Some(working_path), dfs, restore_cfg, pd_cli, Some(pd_control), &br_cx.runtime, truncate_ts.map(|v| v.into_inner()), Some(br_cx.master_key.clone()), reporter, br_cx.object_cache.clone(), limiter, ) }) .await .unwrap()?; let user_view = RestoredKeyspaceView { origin_cluster_id: packed_cluster_id, origin_keyspace_name: packed_keyspace_name, target_cluster_id: req.cluster_id, target_keyspace_name: req.keyspace.clone(), point_in_time: ts_to_datetime(restored.ts), restore_bytes: restored.restore_bytes, tolerated_err: restored.tolerated_err, }; Result::Ok(user_view) }; br_cx.finish_task(req.id, exec.await); } async fn run_create_pack_backup(br: Arc, req: CreatePackBackupRequest) { let cfg = br.config.rl().clone(); let br_cx = &br.context; let mut pcfg = PackConfig::default(); pcfg.dfs = cfg.dfs; pcfg.backup_name = req.backup_name; pcfg.keyspace_name = req.keyspace; pcfg.data_dir = Some(Path::new(&cfg.data_dir).to_owned()); req.s3_override.apply_as_local_override(&mut pcfg.dfs); let reporter = ApiV1xReporter { cx: Arc::clone(br_cx), id: req.id, }; br_cx.with_task(req.id, |task| task.with_mut_state(|s| s.init())); let pcfg2 = pcfg.clone(); let res = async { // NOTE: `create_from_config` cannot be asyncrhonous because it create a runtime // internally... (Somehow it is terrifying to bind a tokio runtime to every S3 // client...) let mut ctx = tokio::task::spawn_blocking(move || PackContext::create_from_config(&pcfg2)) .await .unwrap()?; ctx.set_reporter(Arc::new(reporter)); let (path, meta) = ctx .execute( pcfg.backup_name .strip_suffix(".meta") .unwrap_or(&pcfg.backup_name), ) .await?; let view = PackedBackupView { cluster_id: meta.cluster_id, backup_ts: meta.backup_ts, safe_ts: meta.safe_ts, keyspace_name: meta.get_keyspace_name().to_string(), resolved_ts: meta.resolved_ts, engine_size: meta.engine_size, keyspace_size: meta.keyspace_size, path, }; tokio::task::spawn_blocking(move || drop(ctx)); Result::Ok(view) } .await; br_cx.finish_task(req.id, res); } } #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Task { id: u64, request: TaskRequest, #[serde(default = "Utc::now", skip_deserializing)] create_time: DateTime, #[serde(flatten)] state: TaskState, state_trans: Vec>, } impl Task { pub fn with_mut_state(&mut self, f: impl FnOnce(&mut TaskState)) { let old_state_str = self.state.state_string(); f(&mut self.state); let new_state_str = self.state.state_string(); if new_state_str != old_state_str { info!("API V1X: a task state transition triggers"; "task" => self.id, "from" => old_state_str, "to" => ?self.state); self.state_trans.push(StateTransform { to: new_state_str, at: Utc::now(), }) } } pub fn state(&self) -> &TaskState { &self.state } } #[derive(Deserialize, Clone, Serialize, Debug, PartialEq, Eq)] pub struct CreatePackBackupRequest { #[serde(with = "serde_with::rust::display_fromstr")] id: u64, backup_name: String, keyspace: String, #[serde(flatten)] s3_override: S3Override, } #[derive(Deserialize)] struct TaskRef { #[serde(with = "serde_with::rust::display_fromstr")] id: u64, } #[derive(Deserialize)] pub struct GetBackupRequest { #[serde(flatten)] s3_override: S3Override, #[serde(flatten)] query: GetBackupQuery, } #[derive(Deserialize)] pub struct BackupStatusRequest { start_date: Option, #[serde(flatten)] s3_override: S3Override, } #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] pub struct BackupStatusResponse { pub last_archived_date: Option, pub first_auto_backup_time: Option>, pub last_auto_backup_time: Option>, } #[derive(Deserialize)] #[serde(tag = "query")] pub enum GetBackupQuery { #[serde(rename = "for-pitr")] ForPointInTimeRestore { point_in_time: DateTime }, } #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub struct CreateRestorePackedBackupRequest { #[serde(with = "serde_with::rust::display_fromstr")] id: u64, #[serde(with = "serde_with::rust::display_fromstr")] cluster_id: u64, #[serde(flatten)] exotic: ExoticSource, keyspace: String, #[serde(default)] point_in_time: Option>, /// Max throughput for the restore process, e.g. "500MiB". /// /// - `0` (default): follow the global shared `restore_rate_limit` /// config (if enabled). /// - `>0`: use an independent token bucket for this request (does not /// consume the global bucket). #[serde(default)] max_throughput: ReadableSize, } #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub struct CreateCopyBackupRequest { #[serde(with = "serde_with::rust::display_fromstr")] id: u64, #[serde(flatten)] exotic: ExoticSource, #[serde(flatten)] s3_override: S3Override, /// Copy the backup to the specified prefix. copy_to: String, /// Max throughput for the copy process, e.g. "500MiB". /// /// - `0` (default): follow the global shared `restore_rate_limit` /// config (if enabled). /// - `>0`: use an independent token bucket for this request (does not /// consume the global bucket). #[serde(default)] max_throughput: ReadableSize, } #[derive(Serialize, Debug, Deserialize)] pub struct PackedBackupView { pub cluster_id: u64, pub backup_ts: u64, pub safe_ts: u64, pub keyspace_name: String, pub resolved_ts: u64, pub engine_size: u64, pub keyspace_size: u64, pub path: String, } #[derive(Deserialize)] struct QueryExoticBackupRequest { #[serde(flatten)] exotic: ExoticSource, } #[derive(Serialize, Debug, Deserialize)] pub struct ExoticBackupInfo { #[serde(flatten)] pub base: PackedBackupView, pub backup_time: String, pub backup_size: u64, } #[derive(Serialize, Debug, Deserialize)] pub struct RestoredKeyspaceView { pub origin_cluster_id: u64, pub origin_keyspace_name: String, pub target_cluster_id: u64, pub target_keyspace_name: String, pub point_in_time: DateTime, pub restore_bytes: u64, pub tolerated_err: usize, } #[derive(Serialize, Debug, Clone, Deserialize, PartialEq, Eq)] pub struct S3Override { #[serde( default, rename = "use_bucket", alias = "use_s3_bucket", skip_serializing_if = "Option::is_none" )] pub bucket: Option, #[serde( default, rename = "use_prefix", skip_serializing_if = "Option::is_none" )] pub prefix: Option, /// Redirect the local DFS endpoint for cross-account access. #[serde( default, rename = "use_endpoint", skip_serializing_if = "Option::is_none" )] pub use_endpoint: Option, } /// Cross-account backup source: path + optional endpoint override. #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] pub struct ExoticSource { /// The exotic backup path: `/`. pub exotic_backup: String, /// Azure endpoint for cross-account reads. #[serde(default, skip_serializing_if = "Option::is_none")] pub exotic_backup_endpoint: Option, } /// Extract the account name from an Azure Blob Storage endpoint URL. /// /// Expected format: `https://{account}.blob.core.windows.net` /// Returns the substring between `//` and the first `.`. /// For URLs without a `.` in the host (e.g. `http://localhost:10000`), /// the entire host (including port) is returned. /// Returns `None` if the URL has no `//` or the host part is empty. fn extract_account_name_from_endpoint(ep: &str) -> Option { ep.split("//") .nth(1) .and_then(|s| s.split('.').next()) .filter(|s| !s.is_empty()) .map(|s| s.to_string()) } impl S3Override { /// Fully redirect the local DFS to another account. /// Used for operations that read *and* write entirely within the /// target account (list_backups, pack_backup). fn apply_as_local_override(&self, cfg: &mut DFSConfig) { if let Some(ref ep) = self.use_endpoint { cfg.azure.endpoint = ep.clone(); if let Some(account) = extract_account_name_from_endpoint(ep) { cfg.azure.account_name = account; } // Redirecting to another account requires token-based auth. cfg.azure.account_key.clear(); } if let Some(ref bucket) = self.bucket { cfg.azure.container = bucket.clone(); cfg.s3_bucket = bucket.clone(); } if let Some(ref prefix) = self.prefix { cfg.prefix = prefix.clone(); } } fn enabled(&self) -> bool { self.bucket.is_some() || self.prefix.is_some() || self.use_endpoint.is_some() } } trait ResponseableExt { fn json_with_status(self, status: StatusCode) -> Response; } impl ResponseableExt for HttpResult { fn json_with_status(self, status: StatusCode) -> Response { match self { Ok(v) => make_json_response(status, &v), Err(err_resp) => err_resp.0, } } } impl NativeBrManager { /// Build a DFS that keeps the local bucket/prefix intact but enables /// cross-account reads via exotic_endpoint. /// Use for copy/restore/query_exotic; local redirect fields /// (use_bucket/ use_prefix/use_endpoint) are not accepted — use /// `overridden_storage_local` for list/pack instead. fn overridden_storage(&self, exotic_endpoint: Option<&str>) -> AsyncDropDfs { match exotic_endpoint { Some(ep) if !ep.is_empty() => { let mut cfg = self.config.read().unwrap().dfs.clone(); cfg.azure.exotic_endpoint = ep.to_owned(); kvengine::dfs::new_dfs_from_config(cfg).into() } _ => self.context.dfs.clone().into(), } } /// Build a DFS fully redirected to the target account. /// Use for operations that read and write entirely in the target /// account (list_backups, pack_backup). fn overridden_storage_local(&self, s3_override: &S3Override) -> AsyncDropDfs { if s3_override.enabled() { let mut cfg = self.config.read().unwrap().dfs.clone(); s3_override.apply_as_local_override(&mut cfg); kvengine::dfs::new_dfs_from_config(cfg).into() } else { self.context.dfs.clone().into() } } /// Build a DFS that both redirects local storage AND enables /// cross-account reads via exotic_endpoint. /// Use for copy_backup where we read from exotic and write to /// a redirected local target. fn overridden_storage_full( &self, exotic_endpoint: Option<&str>, s3_override: &S3Override, ) -> AsyncDropDfs { let ep = exotic_endpoint.filter(|e| !e.is_empty()); if ep.is_none() && !s3_override.enabled() { return self.context.dfs.clone().into(); } let mut cfg = self.config.read().unwrap().dfs.clone(); if let Some(ep) = ep { cfg.azure.exotic_endpoint = ep.to_owned(); } if s3_override.enabled() { s3_override.apply_as_local_override(&mut cfg); } kvengine::dfs::new_dfs_from_config(cfg).into() } fn spawn_task(self: &Arc, request: TaskRequest) -> Result { info!("API V1X: about to start a new task."; "request" => ?request); let task_id = request.id(); let mut tasks = self.context.v1x_tasks.wl(); if let Some(task) = tasks.get(&task_id) { // NOTE: maybe add a digest to each task... // So we can know this is a confliction or just a duplicated request. return Err(Error::Existed(format!( "task(id={},req={:?})", task.id, task.request ))); } let state = TaskState::Pending; let task = Task { id: task_id, request: request.clone(), create_time: Utc::now(), state, state_trans: vec![], }; tasks.insert(task_id, task); drop(tasks); self.context.runtime.spawn(request.run(self.clone())); self.context.get_task_for_api(task_id) } } impl BrContext { fn with_task(&self, id: u64, f: impl FnOnce(&mut Task)) { let mut tasks = self.v1x_tasks.wl(); let Some(task) = tasks.get_mut(&id) else { warn!("Reporting error to a removed task."; "task" => %id); return; }; f(task) } fn get_task_for_api(&self, id: u64) -> Result { let tasks = self.v1x_tasks.rl(); let Some(task) = tasks.get(&id) else { return Err(Error::identitied_not_found("task", id)); }; Ok(task.clone()) } fn delete_task_for_api(&self, id: u64) -> Result { let mut tasks = self.v1x_tasks.wl(); let Some(task) = tasks.get(&id) else { return Err(Error::identitied_not_found("task", id)); }; if !matches!(task.state, TaskState::Done { .. } | TaskState::Error { .. }) { return Err(Error::InvalidStateTrans { resource: "v1x::task".to_owned(), from: task.state.state_string(), to: "Absent".to_owned(), }); } Ok(tasks.remove(&id).unwrap()) } fn finish_task(&self, task_id: u64, result: Result) { let mut tasks = self.v1x_tasks.wl(); let Some(task) = tasks.get_mut(&task_id) else { return; }; match result { Ok(val) => { let json_val = serde_json::to_value(val).unwrap(); task.with_mut_state(|s| s.done(json_val)); } Err(err) => task.with_mut_state(|s| s.error(&err)), } } } struct ApiV1xReporter { cx: Arc, id: u64, } impl ApiV1xReporter {} impl ReportPackBackupStepTrait for ApiV1xReporter { fn report_step(&self, step: PackBackupStep) { self.cx.with_task(self.id, |task| { task.with_mut_state(|s| s.running(format!("{step:?}"))) }) } } impl ReportRestoreStepTrait for ApiV1xReporter { fn report_step(&self, step: RestoreStep) { self.cx.with_task(self.id, |task| { task.with_mut_state(|s| s.running(format!("{step:?}"))) }) } } impl ReportUnpackStepTrait for ApiV1xReporter { fn report_step(&self, step: UnpackStep) { self.cx.with_task(self.id, |task| { task.with_mut_state(|s| s.running_info(step.stage_name().to_owned(), step)) }) } } impl ReportCopyStepTrait for ApiV1xReporter { fn report_step(&self, step: native_br::packing::CopyStep) { self.cx.with_task(self.id, |task| { task.with_mut_state(|s| s.running_info(step.stage_name().to_owned(), step)) }) } } fn estimate_packed_backup_size(packed: &PackedBackup) -> u64 { const TABLE_SIZE_MULTIPLIER: f64 = 1.022; let mut total: u64 = 0; for cs in packed.get_shards() { if !cs.has_snapshot() { continue; } let snap = cs.get_snapshot(); for l0 in snap.get_l0_creates() { total += l0.size as u64; } for table in snap.get_table_creates() { total += (table.meta_offset as f64 * TABLE_SIZE_MULTIPLIER) as u64; } for blob in snap.get_blob_creates() { total += blob.meta_offset as u64; } for col in snap.get_columnar_creates() { total += (col.meta_offset as f64 * TABLE_SIZE_MULTIPLIER) as u64; } } total } async fn handle_query_exotic_backup(ctx: HttpRequestContext) -> HttpResult { let req: QueryExoticBackupRequest = ctx.query_params()?; if req.exotic.exotic_backup.is_empty() { return Err(Error::CheckError("exotic_backup path cannot be empty".to_string()).into()); } let dfs = ctx .br .overridden_storage(req.exotic.exotic_backup_endpoint.as_deref()); let env = MigratePackEnv::load_exotic(Arc::clone(&*dfs), &req.exotic.exotic_backup).await?; let meta = env.get_packed_backup(); let backup_size = estimate_packed_backup_size(meta); Ok(ExoticBackupInfo { base: PackedBackupView { cluster_id: meta.cluster_id, backup_ts: meta.backup_ts, safe_ts: meta.safe_ts, keyspace_name: meta.get_keyspace_name().to_string(), resolved_ts: meta.resolved_ts, engine_size: meta.engine_size, keyspace_size: meta.keyspace_size, path: req.exotic.exotic_backup, }, backup_time: ts_to_datetime(meta.backup_ts).to_rfc3339(), backup_size, }) } pub(crate) async fn serve( br_ctx: Arc, req: Request, ) -> hyper::Result> { let ctx = HttpRequestContext::new(req, br_ctx); let path = ctx.raw_req.uri().path().to_owned(); let method = ctx.raw_req.method().clone(); let res = do_serve(ctx).await; if res.status().is_server_error() { warn!("API V1x encounters internal error."; "resp" => ?res, "path" => path, "method" => %method); } Ok(res) } async fn do_serve(mut ctx: HttpRequestContext) -> Response { let path = ctx .raw_req .uri() .path() .split('/') .map(|s| s.to_owned()) .collect::>(); let borrowed_path = path.iter().map(|s| &**s).collect::>(); let method = ctx.raw_req.method().clone(); let not_found = |ctx: &HttpRequestContext| { ( HttpError::error_response( "PathNotFound", StatusCode::NOT_FOUND, json!({"path": ctx.raw_req.uri().path()}), ) .0, "path_not_found", ) }; let method_not_allowed = |ctx: &HttpRequestContext| { ( HttpError::error_response( "MethodNotAllowed", StatusCode::METHOD_NOT_ALLOWED, json!({"path": ctx.raw_req.uri().path(), "method": method.to_string()}), ) .0, "method_not_allowed", ) }; let ob_start_time = tikv_util::time::Instant::now(); info!("API V1X: about to handle a request."; "ctx" => ?ctx); let (res, tag) = match borrowed_path.as_slice() { [.., "task", "copy_backup", id] => match method { Method::PUT => { ctx.request_params.insert("id".to_string(), id.to_string()); let res = handle_copy_pack(ctx) .await .json_with_status(StatusCode::CREATED); (res, "x_create_copy_backup") } _ => method_not_allowed(&ctx), }, [.., "task", "pack_backup", id] => match method { Method::PUT => { ctx.request_params.insert("id".to_string(), id.to_string()); let res = handle_create_pack(ctx) .await .json_with_status(StatusCode::CREATED); (res, "x_create_pack_backup") } _ => method_not_allowed(&ctx), }, [.., "task", "restore_packed", id] => { ctx.request_params.insert("id".to_string(), id.to_string()); match method { Method::PUT => { let res = handle_restore_packed(ctx) .await .json_with_status(StatusCode::CREATED); (res, "x_restore_packed_backup") } _ => method_not_allowed(&ctx), } } [.., "task", id] => { ctx.request_params.insert("id".to_string(), id.to_string()); match method { Method::GET => { let res = handle_get_task(ctx).await.json_with_status(StatusCode::OK); (res, "x_get_task") } Method::DELETE => { let res = handle_delete_task(ctx) .await .json_with_status(StatusCode::OK); (res, "x_delete_task") } _ => method_not_allowed(&ctx), } } [.., "backup"] => match method { Method::GET => { let res = handle_get_backup(ctx) .await .json_with_status(StatusCode::OK); (res, "x_query_backup") } _ => method_not_allowed(&ctx), }, [.., "backup", "status"] => match method { Method::GET => { let res = handle_backup_status(ctx) .await .json_with_status(StatusCode::OK); (res, "x_get_backup_status") } _ => method_not_allowed(&ctx), }, [.., "exotic_backup"] => match method { Method::GET => { let res = handle_query_exotic_backup(ctx) .await .json_with_status(StatusCode::OK); (res, "x_query_exotic_backup") } _ => method_not_allowed(&ctx), }, [.., "backups"] => match method { Method::GET => { let res = handle_list_backup(ctx) .await .json_with_status(StatusCode::OK); (res, "x_list_backup") } _ => method_not_allowed(&ctx), }, [.., "cluster", "flush_wal"] => match method { Method::POST => ( handle_flush_wal(ctx).await.json_with_status(StatusCode::OK), "x_flush_wal", ), _ => method_not_allowed(&ctx), }, [.., "cluster", "lightweight_backup"] => match method { Method::POST => ( handle_create_lightweight_backup(ctx) .await .json_with_status(StatusCode::OK), "x_create_lightweight_backup", ), _ => method_not_allowed(&ctx), }, _ => not_found(&ctx), }; let result = if res.status().is_success() { "success" } else if res.status().is_client_error() { "client_error" } else { "failure" }; NATIVE_BR_HISTOGRAM_VEC .with_label_values(&[tag]) .observe(ob_start_time.saturating_elapsed_secs()); NATIVE_BR_V1X_COUNTER_VEC .with_label_values(&[tag.strip_prefix("x_").unwrap_or(tag), result]) .inc(); res } struct HttpRequestContext { br: Arc, raw_req: Request, request_params: HashMap, } impl Debug for HttpRequestContext { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("HttpRequestContext") .field("method", &self.raw_req.method()) .field("uri", &self.raw_req.uri()) .finish() } } impl HttpRequestContext { fn new(req: Request, br: Arc) -> Self { Self { br, raw_req: req, request_params: Default::default(), } } fn query_params<'de, T: Deserialize<'de>>(&self) -> Result { let query = self.raw_req.uri().query().unwrap_or(""); let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect(); let it = query_pairs .iter() .map(|(k, v)| (k.as_ref(), FieldHinttedStrDe::new(v.as_ref(), k))) .chain( self.request_params .iter() .map(|(k, v)| (k.as_str(), FieldHinttedStrDe::new(v.as_str(), k))), ); let de = MapDeserializer::<_, Error>::new(it); T::deserialize(de) } } async fn handle_get_task(ctx: HttpRequestContext) -> HttpResult { let r: TaskRef = ctx.query_params()?; let task = ctx.br.context.get_task_for_api(r.id)?; Ok(task) } async fn handle_delete_task(ctx: HttpRequestContext) -> HttpResult { let r: TaskRef = ctx.query_params()?; let task = ctx.br.context.delete_task_for_api(r.id)?; Ok(task) } async fn handle_copy_pack(ctx: HttpRequestContext) -> HttpResult { let req: CreateCopyBackupRequest = ctx.query_params()?; ctx.br .spawn_task(TaskRequest::CopyPacked(req.clone())) .map_err(HttpError::from) } async fn handle_create_pack(ctx: HttpRequestContext) -> HttpResult { let req: CreatePackBackupRequest = ctx.query_params()?; ctx.br .spawn_task(TaskRequest::PackBackup(req.clone())) .map_err(HttpError::from) } async fn handle_restore_packed(ctx: HttpRequestContext) -> HttpResult { let req: CreateRestorePackedBackupRequest = ctx.query_params()?; if req.cluster_id != ctx.br.get_cluster_id()? { Err(Error::CheckError(format!( "our cluster is {} but you want to restore {}", ctx.br.get_cluster_id()?, req.cluster_id )))? } let task = ctx.br.spawn_task(TaskRequest::RestorePacked(req.clone()))?; Ok(task) } async fn handle_get_backup(ctx: HttpRequestContext) -> HttpResult { let req: GetBackupRequest = ctx.query_params()?; let GetBackupQuery::ForPointInTimeRestore { point_in_time } = req.query; // When querying a backup contains "15:01:02:600", we need a backup newer than // "150103.meta". let next_second = point_in_time .add(chrono::Duration::milliseconds(500)) .duration_round(chrono::Duration::seconds(1)) .unwrap(); let s3 = ctx.br.overridden_storage_local(&req.s3_override); let (backups, _) = get_all_incremental_backups( s3.as_ref(), &next_second.date_naive(), Some(&next_second.time()), 1, ) .await?; if backups.is_empty() { return Err(Error::noted_not_found( "backup", format_args!( "a suitable backup for point in time restore to {}", point_in_time ), ) .into()); } let backup = backups.into_iter().next().unwrap(); let backup_name = backup_file_full_path(s3.get_prefix(), backup.name().to_owned(), None); let details = fetch_backup_details(s3.as_ref(), &backup_name).await?; let target_ts = TimeStamp::compose(point_in_time.naive_utc().timestamp_millis() as _, 0); if details.safe_ts > target_ts.into_inner() { return Err(Error::noted_not_found("backup", format_args!( "you are querying a backup with timestamp you want({target_ts}); there is one but that timestamp was GCed(to {})", details.safe_ts )).into()); } if details.backup_ts < target_ts.into_inner() { // In case PD local time drifts... // What can we do later in this scenario? Suggest the user use a former // timestamp and try again? return Err(Error::noted_not_found("backup", format_args!( "you are querying a backup with timestamp you want({target_ts}); there is one but that timestamp is beyond the backup time({})", details.backup_ts )).into()); } let actual_backup = Backup { name: backup.name().to_owned(), last_modify_time: backup.created_at().to_rfc3339(), details: Some(details), v1_compat_id: Some(backup.id()), v1_compat_name: Some(backup.created_at().format(BACKUP_NAME_FORMAT).to_string()), }; Ok(actual_backup) } fn earliest_backup_date() -> NaiveDate { NaiveDate::from_ymd_opt(1970, 1, 1).unwrap() } async fn get_first_auto_backup_time( dfs: &dyn Dfs, start_date: NaiveDate, ) -> Result>> { let (backups, _) = get_all_incremental_backups(dfs, &start_date, None, 1).await?; Ok(backups.first().map(|backup| backup.created_at().to_owned())) } async fn get_optional_latest_archive_date( dfs: Arc, start_date: NaiveDate, ) -> Result> { spawn_blocking( move || match get_latest_archive_date(dfs.as_ref(), &start_date) { Ok(date) => Ok(Some(date)), Err(NativeBrError::NoArchiveFoundSince(_)) => Ok(None), Err(err) => Err(err.into()), }, ) .await .unwrap() } async fn get_last_auto_backup_time( dfs: &dyn Dfs, start_date: NaiveDate, ) -> Result>> { let mut start_key = format!("{}/", start_date.format(INCREMENTAL_BACKUP_FOLDER_FORMAT)); let mut last_backup = None; loop { let (files, more, next_start_after) = dfs.list(&start_key, Some("backup/"), None).await?; if let Some(backup) = files .into_iter() .rev() .filter_map(|file| IncrementalBackupFile::try_from_full_path(&file.key)) .next() { last_backup = Some(backup); } if !more { break; } start_key = next_start_after.ok_or_else(|| { Error::Other(box_err!( "dfs list returned has_more without next_start_after" )) })?; } Ok(last_backup.map(|backup| backup.created_at().to_owned())) } async fn handle_backup_status(ctx: HttpRequestContext) -> HttpResult { let req: BackupStatusRequest = ctx.query_params()?; let dfs = ctx.br.overridden_storage_local(&req.s3_override); let start_date = req.start_date.unwrap_or_else(earliest_backup_date); let first_auto_backup_time = get_first_auto_backup_time(dfs.as_ref(), start_date).await?; let last_archived_date = get_optional_latest_archive_date(Arc::clone(&*dfs), start_date).await?; let scan_start_date = last_archived_date.max( first_auto_backup_time .as_ref() .map(|time| time.date_naive()), ); let last_auto_backup_time = match scan_start_date { Some(start_date) => get_last_auto_backup_time(dfs.as_ref(), start_date).await?, None => None, }; Ok(BackupStatusResponse { last_archived_date, first_auto_backup_time, last_auto_backup_time, }) } async fn handle_list_backup(ctx: HttpRequestContext) -> HttpResult { let cluster_id = ctx.br.get_cluster_id()?; let params: ListBackupRequest = ctx.query_params()?; let dfs = ctx.br.overridden_storage_local(¶ms.s3_override); let backup_prefix = "backup/"; let (files, more, _) = dfs .list( ¶ms.from_prefix, Some(backup_prefix), Some(params.max_count as _), ) .await?; let full_prefix = format!("{}/{}", dfs.get_prefix(), backup_prefix); let mut backups = Vec::with_capacity(files.len()); let max_conc = 128; let dfs_ref = dfs.as_ref(); let full_prefix = &full_prefix; let load_details = params.load_details; { let mut st = FuturesUnordered::new(); for v in files { st.push(async move { let v1_compat = v1_compat_info(&v.key); let name = v.key.strip_prefix(full_prefix).unwrap().to_string(); HttpResult::Ok(Backup { name: name.clone(), last_modify_time: v.last_modified, details: if load_details { Some(fetch_backup_details(dfs_ref, &v.key).await?) } else { None }, v1_compat_id: v1_compat.as_ref().map(|(a, _)| *a), v1_compat_name: v1_compat.map(|(_, b)| b), }) }); if st.len() > max_conc { backups.push( st.try_next() .await? .expect("should be something in the stream"), ) } } while let Some(v) = st.try_next().await? { backups.push(v); } } backups.sort_by(|bk1, bk2| bk1.name.cmp(&bk2.name)); Ok(ListBackupResponse { data: backups, cluster_id, has_more: more, }) } const DFS_WORKER_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); #[derive(Deserialize)] struct StoreConfigResponse { rfengine: StoreRfengineConfig, } #[derive(Deserialize)] struct StoreRfengineConfig { #[serde(rename = "lightweight-backup")] lightweight_backup: bool, } async fn store_dfs_worker_is_disabled( store: Store, security_mgr: Arc, client: security::HttpClient, ) -> bool { let disabled = async { let uri = security_mgr .build_uri(format!("{}/config", store.status_address)) .map_err(|err| Error::Other(box_err!("failed to build config URI: {}", err)))?; let request = Request::get(uri).body(Body::empty())?; let (_, body) = send_request_to_store(request, &store, &client, DFS_WORKER_REQUEST_TIMEOUT).await?; let config: StoreConfigResponse = serde_json::from_slice(&body)?; Ok::<_, Error>(!config.rfengine.lightweight_backup) } .await; disabled.unwrap_or_else(|err| { warn!( "failed to get store config when checking DFS worker state"; "store_id" => store.id, "status_address" => &store.status_address, "err" => ?err, ); false }) } pub(super) async fn all_dfs_workers_disabled(pdc: &Arc) -> Result { let stores = crate::get_all_stores_except_tiflash(pdc)?; if stores.is_empty() { return Ok(false); } let security_mgr = pdc.get_security_mgr(); let client = security_mgr .http_client(hyper::Client::builder()) .map_err(|err| Error::Other(box_err!("failed to create http client: {}", err)))?; let disabled: Vec<_> = stores .into_iter() .map(|store| store_dfs_worker_is_disabled(store, security_mgr.clone(), client.clone())) .collect::>() .collect() .await; Ok(disabled.into_iter().all(|disabled| disabled)) } async fn flush_all_wal(pdc: &Arc) -> Result<()> { let mut stores = pdc.get_all_stores(true)?; stores.retain(|store| { let is_tiflash = store .get_labels() .iter() .any(|label| label.get_key() == "engine" && label.get_value().contains("tiflash")); !is_tiflash }); let sec = pdc.get_security_mgr(); let client = sec .http_client(hyper::Client::builder()) .map_err(|e| Error::Other(box_err!("failed to create http client: {}", e)))?; let mut store_requests = vec![]; for store in stores { let sec = &sec; let client = &client; store_requests.push(async move { let req = { let uri = sec .build_uri(format!("{}/rfengine/flush_wal", &store.status_address)) .unwrap(); Request::post(uri).body(Body::empty()) }?; let (status_code, _) = send_request_to_store(req, &store, client, std::time::Duration::from_secs(600)) .await?; if !status_code.is_success() { warn!("backup v1x: flush store failed."; "status_code" => %status_code, "store_id" => store.id); return Err(Error::Other(box_err!( "store {} failed to flush wal; perhaps its dfs worker isn't healthy", store.id ))); } Ok(()) }); } futures::future::try_join_all(store_requests).await?; Ok(()) } async fn handle_create_lightweight_backup( cx: HttpRequestContext, ) -> HttpResult { // "Instant backup" in native-br is a lightweight backup. let timeout = cx.br.config.rl().native_br.instant_backup_timeout.0; let bk = cx .br .context .backup_worker .instant_backup_with_retry(timeout) .await?; flush_all_wal(&cx.br.context.pd_client).await?; Ok(CreateLightweightBackupResponse { backup_name: bk.backup_file.name().to_owned(), }) } async fn handle_flush_wal(cx: HttpRequestContext) -> HttpResult { flush_all_wal(&cx.br.context.pd_client).await?; Ok(FlushResult::default()) } async fn fetch_backup_details(dfs: &dyn Dfs, key: &str) -> Result { let backup_meta_bytes = dfs .get_object(key.to_owned(), key.to_owned(), Default::default()) .await?; let mut backup_meta = ClusterBackupMeta::default(); box_try!(backup_meta.merge_from_bytes(&backup_meta_bytes)); build_backup_details(backup_meta).await } async fn build_backup_details(backup_meta: ClusterBackupMeta) -> Result { let offline_pd = OfflinePd::new(backup_meta.cluster_id, backup_meta.keyspace_meta).await?; let (keyspaces, skipped_keyspace_count) = offline_pd.get_all_keyspaces()?; let mut keyspace_sizes = HashMap::::new(); for store in backup_meta.stores.iter() { for (&keyspace_id, backup_size) in store.get_keyspace_size().iter() { let Some(keyspace_meta) = keyspaces .get(&keyspace_id) .filter(|meta| meta.state == "ENABLED") else { continue; }; let sz = keyspace_sizes .entry(keyspace_id) .or_insert_with(|| KeyspaceBackupInfo { keyspace_id, keyspace_name: keyspace_meta.name.to_owned(), keyspace_state: keyspace_meta.state.to_owned(), size: 0, }); sz.size += backup_size.size; } } Ok(BackupDetails { backup_ts: backup_meta.backup_ts, safe_ts: backup_meta.safe_ts, keyspaces: keyspace_sizes.into_values().collect(), skipped_keyspace_count, }) } #[cfg(test)] mod request_tests { use std::collections::HashMap; use serde::Deserialize; use super::*; fn parse_query(query: &str) -> Result where T: for<'de> Deserialize<'de>, { let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect(); let de = MapDeserializer::<_, Error>::new( query_pairs .iter() .map(|(k, v)| (k.as_ref(), FieldHinttedStrDe::new(v.as_ref(), k))), ); T::deserialize(de) } fn sample_local_override() -> S3Override { S3Override { bucket: Some("target-bucket".to_owned()), prefix: Some("target-prefix".to_owned()), use_endpoint: Some("https://account2.blob.core.windows.net".to_owned()), } } #[tokio::test] async fn test_backup_details_filter_keyspaces_and_count_parse_errors() { let mut backup_meta = ClusterBackupMeta::default(); backup_meta.set_cluster_id(42); for (id, state) in [ (1, kvproto::keyspacepb::KeyspaceState::Enabled), (2, kvproto::keyspacepb::KeyspaceState::Disabled), (3, kvproto::keyspacepb::KeyspaceState::Archived), (4, kvproto::keyspacepb::KeyspaceState::Tombstone), ] { let mut keyspace = kvproto::keyspacepb::KeyspaceMeta::default(); keyspace.set_id(id); keyspace.set_name(format!("ks{id}")); keyspace.set_state(state); backup_meta.mut_keyspace_meta().insert( format!("/pd/42/keyspaces/meta/{id:08}").into_bytes(), keyspace.write_to_bytes().unwrap(), ); } backup_meta .mut_keyspace_meta() .insert(b"/pd/42/keyspaces/meta/00000006".to_vec(), vec![0xff]); let mut negative_timestamp = kvproto::keyspacepb::KeyspaceMeta::default(); negative_timestamp.set_id(7); negative_timestamp.set_name("ks7".to_owned()); negative_timestamp.set_created_at(-1); backup_meta.mut_keyspace_meta().insert( b"/pd/42/keyspaces/meta/00000007".to_vec(), negative_timestamp.write_to_bytes().unwrap(), ); for sizes in [ [(1, 10), (2, 20), (5, 50), (6, 60)], [(1, 11), (3, 30), (4, 40), (7, 70)], ] { let mut store = rfenginepb::StoreBackupMeta::default(); for (id, size) in sizes { let mut backup_size = rfenginepb::BackupSize::default(); backup_size.set_size(size); store.mut_keyspace_size().insert(id, backup_size); } backup_meta.mut_stores().push(store); } let details = build_backup_details(backup_meta).await.unwrap(); assert_eq!(details.keyspaces.len(), 1); let keyspace = &details.keyspaces[0]; assert_eq!(keyspace.keyspace_id, 1); assert_eq!(keyspace.keyspace_name, "ks1"); assert_eq!(keyspace.keyspace_state, "ENABLED"); assert_eq!(keyspace.size, 21); assert_eq!(details.skipped_keyspace_count, 2); } #[test] fn test_requests_accept_canonical_override_names() -> Result<()> { // list and pack use use_endpoint (full local DFS redirect) let local_expected = sample_local_override(); let list_req: ListBackupRequest = parse_query( "from_prefix=backup/&load_details=true&max_count=7\ &use_bucket=target-bucket&use_prefix=target-prefix\ &use_endpoint=https://account2.blob.core.windows.net", )?; assert_eq!(list_req.from_prefix, "backup/"); assert!(list_req.load_details); assert_eq!(list_req.max_count, 7); assert_eq!(list_req.s3_override, local_expected); let get_req: GetBackupRequest = parse_query( "query=for-pitr&point_in_time=2026-03-26T00:00:00Z\ &use_bucket=target-bucket&use_prefix=target-prefix\ &use_endpoint=https://account2.blob.core.windows.net", )?; assert_eq!(get_req.s3_override, local_expected); match get_req.query { GetBackupQuery::ForPointInTimeRestore { point_in_time } => { assert_eq!( point_in_time, DateTime::parse_from_rfc3339("2026-03-26T00:00:00Z") .unwrap() .with_timezone(&Utc) ); } } let pack_req: CreatePackBackupRequest = parse_query( "id=41&backup_name=20260326.meta&keyspace=ks-pack\ &use_bucket=target-bucket&use_prefix=target-prefix\ &use_endpoint=https://account2.blob.core.windows.net", )?; assert_eq!(pack_req.id, 41); assert_eq!(pack_req.backup_name, "20260326.meta"); assert_eq!(pack_req.keyspace, "ks-pack"); assert_eq!(pack_req.s3_override, local_expected); // copy, restore, query: use ExoticSource only; local redirect fields // (use_bucket/use_prefix/use_endpoint) are not part of these requests. let copy_req: CreateCopyBackupRequest = parse_query( "id=42&exotic_backup=src/backup.meta©_to=dst/prefix\ &exotic_backup_endpoint=https://account2.blob.core.windows.net", )?; assert_eq!(copy_req.id, 42); assert_eq!(copy_req.exotic.exotic_backup, "src/backup.meta"); assert_eq!( copy_req.exotic.exotic_backup_endpoint.as_deref(), Some("https://account2.blob.core.windows.net") ); assert_eq!(copy_req.copy_to, "dst/prefix"); let restore_req: CreateRestorePackedBackupRequest = parse_query( "id=43&cluster_id=9527&exotic_backup=src/backup.meta&keyspace=ks-restore\ &point_in_time=2026-03-26T00:00:00Z\ &exotic_backup_endpoint=https://account2.blob.core.windows.net", )?; assert_eq!(restore_req.id, 43); assert_eq!(restore_req.cluster_id, 9527); assert_eq!(restore_req.exotic.exotic_backup, "src/backup.meta"); assert_eq!( restore_req.exotic.exotic_backup_endpoint.as_deref(), Some("https://account2.blob.core.windows.net") ); assert_eq!(restore_req.keyspace, "ks-restore"); assert_eq!( restore_req.point_in_time, Some( DateTime::parse_from_rfc3339("2026-03-26T00:00:00Z") .unwrap() .with_timezone(&Utc) ) ); let query_req: QueryExoticBackupRequest = parse_query( "exotic_backup=src/backup.meta\ &exotic_backup_endpoint=https://account2.blob.core.windows.net", )?; assert_eq!(query_req.exotic.exotic_backup, "src/backup.meta"); assert_eq!( query_req.exotic.exotic_backup_endpoint.as_deref(), Some("https://account2.blob.core.windows.net") ); Ok(()) } // ── apply_as_local_override() behavioral tests ── fn no_override() -> S3Override { S3Override { bucket: None, prefix: None, use_endpoint: None, } } #[test] fn test_enabled_returns_false_when_all_none() { assert!(!no_override().enabled()); } #[test] fn test_enabled_returns_true_for_each_field() { let base = no_override(); assert!( S3Override { bucket: Some("b".into()), ..base.clone() } .enabled() ); assert!( S3Override { prefix: Some("p".into()), ..base.clone() } .enabled() ); assert!( S3Override { use_endpoint: Some("https://ep".into()), ..base } .enabled() ); } #[test] fn test_apply_as_local_override_does_not_set_exotic_endpoint() { let mut cfg = DFSConfig::default(); let s3 = S3Override { use_endpoint: Some("https://local.blob.core.windows.net".to_owned()), bucket: Some("local-bucket".to_owned()), prefix: Some("local-prefix".to_owned()), }; s3.apply_as_local_override(&mut cfg); assert_eq!( cfg.azure.exotic_endpoint, "", "apply_as_local_override must not touch exotic_endpoint" ); // But local fields ARE set. assert_eq!(cfg.azure.endpoint, "https://local.blob.core.windows.net"); assert_eq!(cfg.azure.container, "local-bucket"); assert_eq!(cfg.prefix, "local-prefix"); } #[test] fn test_apply_as_local_override_clears_account_key() { let mut cfg = DFSConfig::default(); cfg.azure.account_key = "secret-key".to_owned(); let s3 = S3Override { use_endpoint: Some("https://account3.blob.core.windows.net".to_owned()), ..no_override() }; s3.apply_as_local_override(&mut cfg); assert_eq!(cfg.azure.account_key, ""); } #[test] fn test_exotic_backup_endpoint_requires_exotic_backup_field() { // Ok type may not implement Debug, so use match instead of unwrap_err(). match parse_query::( "exotic_backup_endpoint=https://exotic.blob.core.windows.net", ) { Ok(_) => panic!("expected missing-field error for QueryExoticBackupRequest"), Err(err) => assert!( err.to_string().contains("missing field"), "unexpected error: {}", err ), } match parse_query::( "id=80©_to=dst\ &exotic_backup_endpoint=https://exotic.blob.core.windows.net", ) { Ok(_) => panic!("expected missing-field error for CreateCopyBackupRequest"), Err(err) => assert!( err.to_string().contains("missing field"), "unexpected error: {}", err ), } match parse_query::( "id=81&cluster_id=1&keyspace=ks\ &exotic_backup_endpoint=https://exotic.blob.core.windows.net", ) { Ok(_) => { panic!("expected missing-field error for CreateRestorePackedBackupRequest") } Err(err) => assert!( err.to_string().contains("missing field"), "unexpected error: {}", err ), } } #[test] fn test_copy_backup_accepts_local_override_params() { let req = parse_query::( "id=42&exotic_backup=src/backup.meta©_to=dst/prefix\ &use_bucket=target-bucket&use_prefix=target-prefix\ &use_endpoint=https://account2.blob.core.windows.net", ) .unwrap(); assert_eq!(req.s3_override.bucket.as_deref(), Some("target-bucket")); assert_eq!(req.s3_override.prefix.as_deref(), Some("target-prefix")); assert_eq!( req.s3_override.use_endpoint.as_deref(), Some("https://account2.blob.core.windows.net") ); } // ── apply_as_local_override tests ──────────────────────────────── #[test] fn test_apply_as_local_override_preserves_connection_string_and_backend() { let mut cfg = DFSConfig::default(); cfg.backend = "azure".to_owned(); cfg.azure.connection_string = "DefaultEndpointsProtocol=https;AccountName=local;AccountKey=key;EndpointSuffix=core.windows.net".to_owned(); cfg.azure.account_key = "original-key".to_owned(); let s3 = S3Override { use_endpoint: Some("https://source.blob.core.windows.net".to_owned()), bucket: Some("source-container".to_owned()), prefix: Some("cse".to_owned()), }; s3.apply_as_local_override(&mut cfg); // backend must remain "azure" so new_dfs_from_config creates AzureFs assert_eq!(cfg.backend, "azure"); // connection_string is NOT cleared by apply_as_local_override. // build_container_client skips it when endpoint is non-empty. assert!(!cfg.azure.connection_string.is_empty()); // account_key must be cleared so redirected access uses MI/token auth assert_eq!(cfg.azure.account_key, ""); // redirected fields ARE set assert_eq!(cfg.azure.endpoint, "https://source.blob.core.windows.net"); assert_eq!(cfg.azure.account_name, "source"); assert_eq!(cfg.azure.container, "source-container"); assert_eq!(cfg.prefix, "cse"); // exotic_endpoint is NOT set by apply_as_local_override assert_eq!(cfg.azure.exotic_endpoint, ""); } #[test] fn test_apply_as_local_override_full_azure_pack_scenario() { let mut cfg = DFSConfig::default(); cfg.backend = "azure".to_owned(); cfg.azure.account_name = "localaccount".to_owned(); cfg.azure.account_key = "local-key".to_owned(); cfg.azure.endpoint = "https://localaccount.blob.core.windows.net".to_owned(); cfg.azure.container = "local-container".to_owned(); cfg.prefix = "local-prefix".to_owned(); let s3 = S3Override { use_endpoint: Some( "https://tidbcloudd6d5ba515313b26.blob.core.windows.net".to_owned(), ), bucket: Some("2034833671956992000-westus3-data".to_owned()), prefix: Some("cse".to_owned()), }; s3.apply_as_local_override(&mut cfg); // After override, the config should point entirely to the source account assert_eq!(cfg.backend, "azure"); assert_eq!( cfg.azure.endpoint, "https://tidbcloudd6d5ba515313b26.blob.core.windows.net" ); assert_eq!(cfg.azure.account_name, "tidbcloudd6d5ba515313b26"); assert_eq!(cfg.azure.account_key, ""); assert_eq!(cfg.azure.container, "2034833671956992000-westus3-data"); assert_eq!(cfg.s3_bucket, "2034833671956992000-westus3-data"); assert_eq!(cfg.prefix, "cse"); // exotic_endpoint must remain empty for pack (read+write in same account) assert_eq!(cfg.azure.exotic_endpoint, ""); } // ── extract_account_name_from_endpoint tests ──────────────────────── #[test] fn test_extract_account_name_standard_azure_url() { assert_eq!( extract_account_name_from_endpoint("https://account2.blob.core.windows.net"), Some("account2".to_owned()) ); } #[test] fn test_extract_account_name_no_dot_in_host() { // localhost:10000 has no '.', so the entire host+port is returned assert_eq!( extract_account_name_from_endpoint("http://localhost:10000"), Some("localhost:10000".to_owned()) ); } #[test] fn test_extract_account_name_azurite_ip() { // Azurite-style IP: first segment before '.' is "127" assert_eq!( extract_account_name_from_endpoint("http://127.0.0.1:10000/devstoreaccount1"), Some("127".to_owned()) ); } #[test] fn test_extract_account_name_empty_string() { assert_eq!(extract_account_name_from_endpoint(""), None); } #[test] fn test_extract_account_name_no_double_slash() { assert_eq!( extract_account_name_from_endpoint("blob.core.windows.net"), None ); } #[test] fn test_extract_account_name_empty_after_double_slash() { assert_eq!(extract_account_name_from_endpoint("https://"), None); assert_eq!( extract_account_name_from_endpoint("https://.blob.core.windows.net"), None ); } } } #[cfg(any(test, feature = "testexport"))] pub mod test_utils { use std::{sync::Arc, time::Duration}; use chrono::{DateTime, NaiveDate, Utc}; use security::{HttpResult, RestfulClient, SecurityManager}; use tikv_util::box_try; use crate::native_br::{ BackupItem, JSON_TIME_FORMAT, ListBackupResponse, RestoreProgressResponse, v1x::{CreateLightweightBackupResponse, FlushResult}, }; #[derive(Default, Serialize, Deserialize, Debug)] #[serde(default)] struct DummyRequest {} fn append_s3_override( ser: &mut url::form_urlencoded::Serializer<'_, String>, s3: &super::v1x::S3Override, ) { if let Some(bucket) = s3.bucket.as_deref() { ser.append_pair("use_bucket", bucket); } if let Some(prefix) = s3.prefix.as_deref() { ser.append_pair("use_prefix", prefix); } if let Some(ep) = s3.use_endpoint.as_deref() { ser.append_pair("use_endpoint", ep); } } fn append_s3_override_legacy( ser: &mut url::form_urlencoded::Serializer<'_, String>, s3: &super::v1x::S3Override, ) { if let Some(bucket) = s3.bucket.as_deref() { ser.append_pair("use_s3_bucket", bucket); } if let Some(prefix) = s3.prefix.as_deref() { ser.append_pair("use_prefix", prefix); } if let Some(ep) = s3.use_endpoint.as_deref() { ser.append_pair("use_endpoint", ep); } } pub struct NativeBrSvcClient { cluster_id: u64, inner: RestfulClient, } impl NativeBrSvcClient { pub fn new( cluster_id: u64, endpoints: Vec, security_mgr: Arc, ) -> HttpResult { Ok(Self { cluster_id, inner: RestfulClient::new("native_br_cli", endpoints, security_mgr)?, }) } /// Sets the timeout used by each HTTP request and by endpoint retries. pub fn set_timeout(&mut self, request_timeout: Duration, retry_timeout: Duration) { self.inner.set_timeout(request_timeout, retry_timeout); } pub async fn list_backups( &self, last_backup_time: &DateTime, ) -> HttpResult { let query = url::form_urlencoded::Serializer::new(String::new()) .extend_pairs([ ("cluster_id", self.cluster_id.to_string()), ( "last_backup_time", last_backup_time.format(JSON_TIME_FORMAT).to_string(), ), ]) .finish(); self.inner.get(format!("api/v1/backups?{query}")).await } pub async fn list_backups_x( &self, from_prefix: String, load_details: bool, max_count: u64, ) -> HttpResult { let query = url::form_urlencoded::Serializer::new(String::new()) .extend_pairs([ ("from_prefix", from_prefix), ("load_details", load_details.to_string()), ("max_count", max_count.to_string()), ]) .finish(); Ok(box_try!( self.inner.get(format!("api/v1/x/backups?{query}")).await )) } pub async fn restore_keyspace_to_backup( &self, restore_id: u64, keyspace: String, source_keyspace: Option, backup: &BackupItem, ) -> HttpResult { self.restore_keyspace(restore_id, keyspace, source_keyspace, backup) .await } pub async fn pack_backup( &self, id: u64, backup_name: &str, keyspace_name: &str, ) -> HttpResult { // Backward compatible: call the variant with no S3 override. self.pack_backup_with_override( id, backup_name, keyspace_name, super::v1x::S3Override { bucket: None, prefix: None, use_endpoint: None, }, ) .await } pub async fn pack_backup_with_override( &self, id: u64, backup_name: &str, keyspace_name: &str, s3_override: super::v1x::S3Override, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); ser.append_pair("backup_name", backup_name); ser.append_pair("keyspace", keyspace_name); append_s3_override(&mut ser, &s3_override); let query = ser.finish(); Ok(box_try!( self.inner .put( format!("api/v1/x/task/pack_backup/{id}?{query}"), &DummyRequest {} ) .await )) } pub async fn pack_backup_with_legacy_override( &self, id: u64, backup_name: &str, keyspace_name: &str, s3_override: super::v1x::S3Override, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); ser.append_pair("backup_name", backup_name); ser.append_pair("keyspace", keyspace_name); append_s3_override_legacy(&mut ser, &s3_override); let query = ser.finish(); Ok(box_try!( self.inner .put( format!("api/v1/x/task/pack_backup/{id}?{query}"), &DummyRequest {} ) .await )) } pub async fn restore_packed_backup( &self, id: u64, exotic_path: &str, keyspace_name: &str, point_in_time: Option>, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); ser.append_pair("cluster_id", &self.cluster_id.to_string()); ser.append_pair("exotic_backup", exotic_path); ser.append_pair("keyspace", keyspace_name); if let Some(pt) = point_in_time { ser.append_pair("point_in_time", &pt.to_rfc3339()); } let query = ser.finish(); Ok(box_try!( self.inner .put( format!("api/v1/x/task/restore_packed/{id}?{query}"), &DummyRequest {} ) .await )) } pub async fn copy_packed_backup( &self, id: u64, exotic_path: &str, copy_to: &str, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); ser.append_pair("exotic_backup", exotic_path); ser.append_pair("copy_to", copy_to); let query = ser.finish(); Ok(box_try!( self.inner .put( format!("api/v1/x/task/copy_backup/{id}?{query}"), &DummyRequest {} ) .await )) } // Query a suitable backup for point-in-time restore via API v1x. pub async fn get_backup_for_pitr( &self, point_in_time: DateTime, ) -> HttpResult { self.get_backup_for_pitr_with_override( point_in_time, super::v1x::S3Override { bucket: None, prefix: None, use_endpoint: None, }, ) .await } // Variant that allows overriding the target S3 bucket/prefix. pub async fn get_backup_for_pitr_with_override( &self, point_in_time: DateTime, s3_override: super::v1x::S3Override, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); // match GetBackupQuery::ForPointInTimeRestore { .. } ser.append_pair("query", "for-pitr"); ser.append_pair("point_in_time", &point_in_time.to_rfc3339()); append_s3_override(&mut ser, &s3_override); let query = ser.finish(); Ok(box_try!( self.inner.get(format!("api/v1/x/backup?{query}")).await )) } pub async fn get_backup_for_pitr_with_legacy_override( &self, point_in_time: DateTime, s3_override: super::v1x::S3Override, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); ser.append_pair("query", "for-pitr"); ser.append_pair("point_in_time", &point_in_time.to_rfc3339()); append_s3_override_legacy(&mut ser, &s3_override); let query = ser.finish(); Ok(box_try!( self.inner.get(format!("api/v1/x/backup?{query}")).await )) } pub async fn get_backup_status( &self, start_date: Option, ) -> HttpResult { let mut ser = url::form_urlencoded::Serializer::new(String::new()); if let Some(start_date) = start_date { ser.append_pair("start_date", &start_date.to_string()); } let query = ser.finish(); Ok(box_try!( self.inner .get(format!("api/v1/x/backup/status?{query}")) .await )) } pub async fn get_task_status(&self, id: u64) -> HttpResult { Ok(box_try!( self.inner.get(format!("api/v1/x/task/{id}")).await )) } pub async fn restore_keyspace( &self, restore_id: u64, keyspace: String, source_keyspace: Option, backup: &BackupItem, ) -> HttpResult { let mut serializer = url::form_urlencoded::Serializer::new(String::new()); serializer.extend_pairs([ ("cluster_id", self.cluster_id.to_string()), ("keyspace", keyspace), ("backup_id", backup.id.to_string()), ("backup_name", backup.name.clone()), ]); if let Some(source_keyspace) = source_keyspace { serializer.append_pair("source_keyspace", &source_keyspace); } let query = serializer.finish(); self.inner .put( format!("api/v1/restore_keyspace/{restore_id}?{query}"), &DummyRequest {}, ) .await } pub async fn restore_keyspace_to_point_in_time( &self, restore_id: u64, keyspace: String, point_in_time: DateTime, ) -> HttpResult { let query = url::form_urlencoded::Serializer::new(String::new()) .extend_pairs([ ("cluster_id", self.cluster_id.to_string()), ("keyspace", keyspace), ( "point_in_time", point_in_time.format(JSON_TIME_FORMAT).to_string(), ), ]) .finish(); self.inner .put( format!("api/v1/restore_keyspace/{restore_id}?{query}"), &DummyRequest {}, ) .await } pub async fn restore_keyspace_from_online_snapshot( &self, restore_id: u64, source_keyspace: String, target_keyspace: String, ) -> HttpResult { let query = url::form_urlencoded::Serializer::new(String::new()) .extend_pairs([ ("cluster_id", self.cluster_id.to_string()), ("keyspace", target_keyspace), ("source_keyspace", source_keyspace), ("restore_mode", "online_snapshot".to_string()), ]) .finish(); self.inner .put( format!("api/v1/restore_keyspace/{restore_id}?{query}"), &DummyRequest {}, ) .await } pub async fn get_restore_progress( &self, restore_id: u64, keyspace: String, ) -> HttpResult { let query = url::form_urlencoded::Serializer::new(String::new()) .extend_pairs([ ("cluster_id", self.cluster_id.to_string()), ("keyspace", keyspace), ]) .finish(); self.inner .get(format!("api/v1/restore_keyspace/{restore_id}?{query}")) .await } pub fn flush_wals(&self) -> HttpResult { futures::executor::block_on(async { Ok(box_try!( self.inner .post("api/v1/x/cluster/flush_wal", &DummyRequest {}) .await )) }) } pub fn create_lightweight_backup(&self) -> HttpResult { futures::executor::block_on(async { Ok(box_try!( self.inner .post("api/v1/x/cluster/lightweight_backup", &DummyRequest {}) .await )) }) } /// Creates a lightweight backup without blocking the calling async /// executor. pub async fn create_lightweight_backup_async( &self, ) -> HttpResult { Ok(box_try!( self.inner .post("api/v1/x/cluster/lightweight_backup", &DummyRequest {}) .await )) } } } #[cfg(test)] mod tests { use std::str::FromStr; use kvengine::dfs::test_util::new_test_s3fs; use pd_client::PdClient; use security::GetSecurityManager; use tikv_util::config::ReadableDuration; use super::*; use crate::native_br::RestoreState::{Init, Running, Succeed}; #[test] fn test_restore_task_state() { let (mgr, _temp_dir) = new_test_br_manager("2s"); let ctx = &mgr.context; // TODO: test for illegal state transition. let states_cases = vec![ vec![Init], vec![Init, Running], vec![Init, Running, Succeed], vec![Init, Running, RestoreState::Error], vec![Init, Running, RestoreState::Error, Init], ]; let backup = IncrementalBackupFile::from_datetime(Utc::now()); let restore_params = RestoreParams { backup_name: Some(backup.name().to_string()), ..Default::default() }; for (restore_id, states) in states_cases.into_iter().enumerate() { for state in states { let res = if state == Init { ctx.init_restore_state( restore_id as u64, &format!("ks{}", restore_id), RestoreType::Normal, restore_params.clone(), ) } else { ctx.change_restore_state( restore_id as u64, &format!("ks{}", restore_id), state, None, 0, 0, 0, ) }; assert!(res.unwrap()); } } assert_eq!(mgr.get_not_final_task_count(), 3); assert_eq!(mgr.get_all_restore_task().len(), 5); thread::sleep(Duration::from_secs(2)); mgr.cleanup_expired_restores(); assert_eq!(mgr.get_not_final_task_count(), 3); assert_eq!(mgr.get_all_restore_task().len(), 3); assert!( ctx.change_restore_state(1, "ks1", Succeed, None, 0, 0, 0) .unwrap() ); mgr.cleanup_expired_restores(); assert_eq!(mgr.get_not_final_task_count(), 2); assert_eq!(mgr.get_all_restore_task().len(), 3); // Test restore_params checking. // Incorrect keyspace name. ctx.change_restore_state(0, "ks1", Running, None, 0, 0, 0) .unwrap_err(); // Transit to running. assert!( ctx.change_restore_state(0, "ks0", Running, None, 0, 0, 0) .unwrap() ); assert!( ctx.change_restore_state(0, "ks0", RestoreState::Error, None, 0, 0, 0) .unwrap() ); // Incorrect backup. let mut params1 = restore_params.clone(); let backup1 = IncrementalBackupFile::try_from_full_path("cse-local/backup/20230321/112233.meta") .unwrap(); params1.backup_name = Some(backup1.name().to_string()); ctx.init_restore_state(0, "ks0", RestoreType::Normal, params1) .unwrap_err(); } #[test] fn test_restore_task_meta() { // Test compatibility with meta of old version (without "restore_params"). let mut meta = serde_json::from_str::( r#"{"keyspace_name":"ks","restore_type":"Normal","start":1000}"#, ) .unwrap(); assert!(meta.restore_params.is_none()); for params in [ RestoreParams::new( &RestoreSource::Normal(IncrementalBackupFile::from_datetime(Utc::now())), "ks0", "ks1", ) .unwrap(), RestoreParams::new( &RestoreSource::Pitr { backup: None, point_in_time: Utc::now(), }, "ks2", "ks2", ) .unwrap(), RestoreParams::new(&RestoreSource::OnlineSnapshot, "ks3", "ks4").unwrap(), ] { meta.restore_params = params; assert_eq!( serde_json::from_slice::(&serde_json::to_vec(&meta).unwrap()) .unwrap(), meta ); } } #[test] fn test_restore_params() { let utc = Utc::now(); let backup = IncrementalBackupFile::from_datetime(utc); assert_eq!( RestoreParams::new(&RestoreSource::Normal(backup.clone()), "ks0", "ks0",).unwrap(), RestoreParams { backup_name: Some(backup.name().to_string()), ..Default::default() } ); assert_eq!( RestoreParams::new( &RestoreSource::Pitr { backup: Some(backup), point_in_time: utc, }, "ks1", "ks0", ) .unwrap(), RestoreParams { point_in_time: Some(utc.timestamp()), source_keyspace: Some("ks1".to_string()), ..Default::default() } ); assert_eq!( RestoreParams::new( &RestoreSource::Pitr { backup: None, point_in_time: utc, }, "ks1", "ks1", ) .unwrap(), RestoreParams { point_in_time: Some(utc.timestamp()), ..Default::default() } ); assert_eq!( RestoreParams::new(&RestoreSource::OnlineSnapshot, "ks1", "ks2",).unwrap(), RestoreParams { source_keyspace: Some("ks1".to_string()), ..Default::default() } ); RestoreParams::new(&RestoreSource::OnlineSnapshot, "ks1", "ks1").unwrap_err(); } #[tokio::test] async fn test_parse_online_snapshot_restore_source() { async fn parse(manager: &NativeBrManager, query: &str) -> Result { let pairs = url::form_urlencoded::parse(query.as_bytes()).collect(); parse_restore_source(manager, &pairs).await } let (manager, _dir) = new_test_br_manager("1h"); assert!(matches!( parse(&manager, "restore_mode=online_snapshot") .await .unwrap(), RestoreSource::OnlineSnapshot )); parse(&manager, "restore_mode=online_snapshot&backup_id=1") .await .unwrap_err(); parse(&manager, "restore_mode=online_snapshot&point_in_time=now") .await .unwrap_err(); parse(&manager, "restore_mode=unknown").await.unwrap_err(); } #[test] fn test_native_br_manager_stores_master_key_in_context() { let thread_pool = Arc::new(AsyncSafeRuntime::new( tokio::runtime::Builder::new_multi_thread() .enable_all() .worker_threads(1) .thread_name("br_manager_test") .build() .unwrap(), )); let file_data = "abcdefgh".to_string().into_bytes(); let dfs = Arc::new(new_test_s3fs(&file_data)); let temp_dir = tempfile::tempdir().unwrap(); let master_key = MasterKey::new(&[9; 32]); let manager = NativeBrManager::new( thread_pool, Arc::new(MockPdClient {}), dfs, temp_dir.path().to_path_buf(), master_key.clone(), Config::default(), ); assert_eq!(manager.context.master_key.export(), master_key.export()); } fn new_test_br_manager(restore_task_ttl: &str) -> (NativeBrManager, tempfile::TempDir) { let thread_pool = Arc::new(AsyncSafeRuntime::new( tokio::runtime::Builder::new_multi_thread() .enable_all() .worker_threads(1) .thread_name("br_manager_test") .build() .unwrap(), )); let file_data = "abcdefgh".to_string().into_bytes(); let dfs = Arc::new(new_test_s3fs(&file_data)); let temp_dir = tempfile::tempdir().unwrap(); ( NativeBrManager::new( thread_pool, Arc::new(MockPdClient {}), dfs, temp_dir.path().to_path_buf(), MasterKey::new(&[1; 32]), Config { native_br: NativeBrConfig { restore_task_ttl: ReadableDuration::from_str(restore_task_ttl).unwrap(), ..Default::default() }, ..Default::default() }, ), temp_dir, ) } struct MockPdClient {} impl PdClient for MockPdClient {} impl GetSecurityManager for MockPdClient {} }