// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    borrow::Cow,
    cell::RefCell,
    fmt,
    sync::{Arc, Mutex, RwLock, atomic::*, mpsc},
    time::{SystemTime, UNIX_EPOCH},
};

use async_channel::SendError;
use concurrency_manager::ConcurrencyManager;
use engine_traits::{CfName, SstCompressionType, name_to_cf};
use external_storage::{BackendConfig, HdfsConfig};
use external_storage_export::{ExternalStorage, create_storage};
use futures::channel::mpsc::*;
use kvproto::{
    brpb::*,
    encryptionpb::EncryptionMethod,
    kvrpcpb::{ApiVersion, Context, IsolationLevel, KeyRange},
    metapb::*,
};
use online_config::OnlineConfig;
use raftstore::coprocessor::RegionInfoProvider;
use tikv::{
    config::BackupConfig,
    storage::{
        Statistics,
        kv::{Engine, SnapContext},
        mvcc::Error as MvccError,
        txn::{CloudStore, EntryBatch, Error as TxnError, TxnEntryScanner, TxnEntryStore},
    },
};
use tikv_util::{
    box_err, debug, error, error_unknown, impl_display_as_debug, info,
    store::find_peer,
    time::{Instant, Limiter},
    warn,
    worker::Runnable,
};
use tokio::runtime::Runtime;
use txn_types::{Key, Lock, TimeStamp};

use crate::{
    Error,
    metrics::*,
    softlimit::{CpuStatistics, SoftLimit, SoftLimitByCpu},
    utils::{ControlThreadPool, KeyValueCodec},
    writer::BackupWriterBuilder,
    *,
};

const BACKUP_BATCH_LIMIT: usize = 1024;

#[derive(Clone)]
struct Request {
    start_key: Vec<u8>,
    end_key: Vec<u8>,
    sub_ranges: Vec<KeyRange>,
    start_ts: TimeStamp,
    end_ts: TimeStamp,
    limiter: Limiter,
    backend: StorageBackend,
    cancel: Arc<AtomicBool>,
    dst_api_ver: ApiVersion,
    cf: CfName,
    compression_type: CompressionType,
    compression_level: i32,
    cipher: CipherInfo,
}

/// Backup Task.
pub struct Task {
    request: Request,
    pub(crate) resp: UnboundedSender<BackupResponse>,
}

impl_display_as_debug!(Task);

impl fmt::Debug for Task {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("BackupTask")
            .field("start_ts", &self.request.start_ts)
            .field("end_ts", &self.request.end_ts)
            .field(
                "start_key",
                &log_wrappers::Value::key(&self.request.start_key),
            )
            .field("end_key", &log_wrappers::Value::key(&self.request.end_key))
            .field("dst_api_ver", &self.request.dst_api_ver)
            .field("cf", &self.request.cf)
            .finish()
    }
}

impl Task {
    /// Create a backup task based on the given backup request.
    pub fn new(
        req: BackupRequest,
        resp: UnboundedSender<BackupResponse>,
    ) -> Result<(Task, Arc<AtomicBool>)> {
        let cancel = Arc::new(AtomicBool::new(false));

        let speed_limit = req.get_rate_limit();
        let limiter = Limiter::new(if speed_limit > 0 {
            speed_limit as f64
        } else {
            f64::INFINITY
        });
        let cf = name_to_cf(req.get_cf()).ok_or_else(|| crate::Error::InvalidCf {
            cf: req.get_cf().to_owned(),
        })?;

        let task = Task {
            request: Request {
                start_key: req.get_start_key().to_owned(),
                end_key: req.get_end_key().to_owned(),
                sub_ranges: req.get_sub_ranges().to_owned(),
                start_ts: req.get_start_version().into(),
                end_ts: req.get_end_version().into(),
                backend: req.get_storage_backend().clone(),
                limiter,
                cancel: cancel.clone(),
                dst_api_ver: req.get_dst_api_version(),
                cf,
                compression_type: req.get_compression_type(),
                compression_level: req.get_compression_level(),
                cipher: req.cipher_info.unwrap_or_else(|| {
                    let mut cipher = CipherInfo::default();
                    cipher.set_cipher_type(EncryptionMethod::Plaintext);
                    cipher
                }),
            },
            resp,
        };
        Ok((task, cancel))
    }

    /// Check whether the task is canceled.
    pub fn has_canceled(&self) -> bool {
        self.request.cancel.load(Ordering::SeqCst)
    }
}

#[derive(Debug)]
pub struct BackupRange {
    start_key: Option<Key>,
    end_key: Option<Key>,
    region: Region,
    leader: Peer,
}

/// The generic saveable writer. for generic `InMemBackupFiles`.
/// Maybe what we really need is make Writer a trait...
enum KvWriter {
    Txn(BackupWriter),
}

impl fmt::Debug for KvWriter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Txn(_) => f.debug_tuple("Txn").finish(),
        }
    }
}

impl KvWriter {
    async fn save(self, storage: &dyn ExternalStorage) -> Result<Vec<File>> {
        match self {
            Self::Txn(writer) => writer.save(storage).await,
        }
    }

    fn need_flush_keys(&self) -> bool {
        match self {
            Self::Txn(writer) => writer.need_flush_keys(),
        }
    }
}

#[derive(Debug)]
struct InMemBackupFiles {
    files: KvWriter,
    start_key: Vec<u8>,
    end_key: Vec<u8>,
    start_version: TimeStamp,
    end_version: TimeStamp,
    region: Region,
}

async fn save_backup_file_worker(
    rx: async_channel::Receiver<InMemBackupFiles>,
    tx: UnboundedSender<BackupResponse>,
    storage: Arc<dyn ExternalStorage>,
    codec: KeyValueCodec,
) {
    while let Ok(msg) = rx.recv().await {
        let files = if msg.files.need_flush_keys() {
            match msg.files.save(&storage).await {
                Ok(mut split_files) => {
                    for file in split_files.iter_mut() {
                        file.set_start_key(msg.start_key.clone());
                        file.set_end_key(msg.end_key.clone());
                        file.set_start_version(msg.start_version.into_inner());
                        file.set_end_version(msg.end_version.into_inner());
                    }
                    Ok(split_files)
                }
                Err(e) => {
                    error_unknown!(?e; "backup save file failed");
                    Err(e)
                }
            }
        } else {
            Ok(vec![])
        };
        let mut response = BackupResponse::default();
        match files {
            Err(e) => {
                error_unknown!(?e; "backup region failed";
                    "region" => ?msg.region,
                    "start_key" => &log_wrappers::Value::key(&msg.start_key),
                    "end_key" => &log_wrappers::Value::key(&msg.end_key),
                );
                response.set_error(e.into());
            }
            Ok(files) => {
                response.set_files(files.into());
            }
        }
        response.set_start_key(msg.start_key.clone());
        response.set_end_key(msg.end_key.clone());
        response.set_api_version(codec.dst_api_ver);
        if let Err(e) = tx.unbounded_send(response) {
            error_unknown!(?e; "backup failed to send response"; "region" => ?msg.region,
            "start_key" => &log_wrappers::Value::key(&msg.start_key),
            "end_key" => &log_wrappers::Value::key(&msg.end_key),);
            if e.is_disconnected() {
                return;
            }
        }
    }
}

/// Send the save task to the save worker.
/// Record the wait time at the same time.
async fn send_to_worker_with_metrics(
    tx: &async_channel::Sender<InMemBackupFiles>,
    files: InMemBackupFiles,
) -> std::result::Result<(), SendError<InMemBackupFiles>> {
    let files = match tx.try_send(files) {
        Ok(_) => return Ok(()),
        Err(e) => e.into_inner(),
    };
    let begin = Instant::now();
    tx.send(files).await?;
    BACKUP_SCAN_WAIT_FOR_WRITER_HISTOGRAM.observe(begin.saturating_elapsed_secs());
    Ok(())
}

impl BackupRange {
    /// Get entries from the scanner and save them to storage
    async fn backup<E: Engine>(
        &self,
        writer_builder: BackupWriterBuilder,
        mut engine: E,
        concurrency_manager: ConcurrencyManager,
        backup_ts: TimeStamp,
        begin_ts: TimeStamp,
        saver: async_channel::Sender<InMemBackupFiles>,
        storage_name: &str,
    ) -> Result<Statistics> {
        let mut ctx = Context::default();
        ctx.set_region_id(self.region.get_id());
        ctx.set_region_epoch(self.region.get_region_epoch().to_owned());
        ctx.set_peer(self.leader.clone());

        // Update max_ts and check the in-memory lock table before getting the snapshot
        concurrency_manager
            .update_max_ts(backup_ts, "backup_range")
            .map_err(TxnError::from)?;
        concurrency_manager
            .read_range_check(
                self.start_key.as_ref(),
                self.end_key.as_ref(),
                |key, lock| {
                    Lock::check_ts_conflict(
                        Cow::Borrowed(lock),
                        key,
                        backup_ts,
                        &Default::default(),
                        IsolationLevel::Si,
                    )
                },
            )
            .map_err(MvccError::from)
            .map_err(TxnError::from)?;

        // Currently backup always happens on the leader, so we don't need
        // to set key ranges and start ts to check.
        assert!(!ctx.get_replica_read());
        let snap_ctx = SnapContext {
            pb_ctx: &ctx,
            ..Default::default()
        };

        let start_snapshot = Instant::now();
        let snapshot = match engine.snapshot(snap_ctx) {
            Ok(s) => s,
            Err(e) => {
                error!(?e; "backup snapshot failed");
                return Err(e.into());
            }
        };
        BACKUP_RANGE_HISTOGRAM_VEC
            .with_label_values(&["snapshot"])
            .observe(start_snapshot.saturating_elapsed().as_secs_f64());
        let cloud_store = CloudStore::new(
            snapshot,
            backup_ts.into_inner(),
            Default::default(),
            false,
            false,
        );
        let start_key = self.start_key.clone();
        let end_key = self.end_key.clone();
        // Incremental backup needs to output delete records.
        let incremental = !begin_ts.is_zero();
        let mut scanner = cloud_store
            .entry_scanner(start_key, end_key, begin_ts, incremental)
            .unwrap();

        let start_scan = Instant::now();
        let mut batch = EntryBatch::with_capacity(BACKUP_BATCH_LIMIT);
        let mut next_file_start_key = self
            .start_key
            .clone()
            .map_or_else(Vec::new, |k| k.into_raw().unwrap());
        let mut writer = writer_builder.build(next_file_start_key.clone(), storage_name)?;
        loop {
            if let Err(e) = scanner.scan_entries(&mut batch) {
                error!(?e; "backup scan entries failed");
                return Err(e.into());
            };
            if batch.is_empty() {
                break;
            }
            debug!("backup scan entries"; "len" => batch.len());

            let entries = batch.drain();
            if writer.need_split_keys() {
                let this_end_key = entries.as_slice().first().map_or_else(
                    || Err(Error::Other(box_err!("get entry error: nothing in batch"))),
                    |x| {
                        x.to_key().map(|k| k.into_raw().unwrap()).map_err(|e| {
                            error!(?e; "backup save file failed");
                            Error::Other(box_err!("Decode error: {:?}", e))
                        })
                    },
                )?;
                let this_start_key = next_file_start_key.clone();
                let msg = InMemBackupFiles {
                    files: KvWriter::Txn(writer),
                    start_key: this_start_key,
                    end_key: this_end_key.clone(),
                    start_version: begin_ts,
                    end_version: backup_ts,
                    region: self.region.clone(),
                };
                send_to_worker_with_metrics(&saver, msg).await?;
                next_file_start_key = this_end_key;
                writer = writer_builder
                    .build(next_file_start_key.clone(), storage_name)
                    .map_err(|e| {
                        error_unknown!(?e; "backup writer failed");
                        e
                    })?;
            }

            // Build sst files.
            if let Err(e) = writer.write(entries, true) {
                error_unknown!(?e; "backup build sst failed");
                return Err(e);
            }
        }
        drop(cloud_store);
        let stat = scanner.take_statistics();
        let take = start_scan.saturating_elapsed_secs();
        if take > 30.0 {
            warn!("backup scan takes long time.";
                "take(s)" => %take,
                "region" => ?self.region,
                "start_key" => %redact_option_key(&self.start_key),
                "end_key" => %redact_option_key(&self.end_key),
                "stat" => ?stat,
            );
        }
        BACKUP_RANGE_HISTOGRAM_VEC
            .with_label_values(&["scan"])
            .observe(take);

        let msg = InMemBackupFiles {
            files: KvWriter::Txn(writer),
            start_key: next_file_start_key,
            end_key: self
                .end_key
                .clone()
                .map_or_else(Vec::new, |k| k.into_raw().unwrap()),
            start_version: begin_ts,
            end_version: backup_ts,
            region: self.region.clone(),
        };
        send_to_worker_with_metrics(&saver, msg).await?;

        Ok(stat)
    }
}

#[derive(Clone)]
pub struct ConfigManager(Arc<RwLock<BackupConfig>>);

impl online_config::ConfigManager for ConfigManager {
    fn dispatch(&mut self, change: online_config::ConfigChange) -> online_config::Result<()> {
        self.0.write().unwrap().update(change)
    }
}

#[derive(Clone)]
struct SoftLimitKeeper {
    limit: SoftLimit,
    config: ConfigManager,
}

impl SoftLimitKeeper {
    fn new(config: ConfigManager) -> Self {
        let BackupConfig { num_threads, .. } = *config.0.read().unwrap();
        let limit = SoftLimit::new(num_threads);
        BACKUP_SOFTLIMIT_GAUGE.set(limit.current_cap() as _);
        Self { limit, config }
    }

    /// run the soft limit keeper at background.
    async fn run(self) {
        let mut cpu_quota =
            SoftLimitByCpu::with_remain(self.config.0.read().unwrap().auto_tune_remain_threads);
        loop {
            if let Err(err) = self.on_tick(&mut cpu_quota).await {
                warn!("soft limit on_tick failed."; "err" => %err);
            }
            let auto_tune_refresh_interval =
                self.config.0.read().unwrap().auto_tune_refresh_interval;
            tokio::time::sleep(auto_tune_refresh_interval.0).await;
        }
    }

    async fn on_tick<S: CpuStatistics>(&self, cpu_quota: &mut SoftLimitByCpu<S>) -> Result<()> {
        let BackupConfig {
            enable_auto_tune,
            num_threads,
            auto_tune_remain_threads,
            ..
        } = *self.config.0.read().unwrap();
        cpu_quota.set_remain(auto_tune_remain_threads);

        let mut quota_val = num_threads;
        if enable_auto_tune {
            quota_val = cpu_quota
                .get_quota(|s| s.contains("bkwkr"))
                .clamp(1, num_threads);
        }

        self.limit.resize(quota_val).await.map_err(|err| {
            warn!(
                "error during appling the soft limit for backup.";
                "current_limit" => %self.limit.current_cap(),
                "to_set_value" => %quota_val,
                "err" => %err,
            );
            Error::Other(box_err!("failed to resize softlimit: {}", err))
        })?;

        BACKUP_SOFTLIMIT_GAUGE.set(self.limit.current_cap() as _);
        Ok(())
    }

    fn limit(&self) -> SoftLimit {
        self.limit.clone()
    }
}

/// The endpoint of backup.
///
/// It coordinates backup tasks and dispatches them to different workers.
pub struct Endpoint<E: Engine, R: RegionInfoProvider + Clone + 'static> {
    store_id: u64,
    pool: RefCell<ControlThreadPool>,
    io_pool: Runtime,
    config_manager: ConfigManager,
    concurrency_manager: ConcurrencyManager,
    softlimit: SoftLimitKeeper,
    api_version: ApiVersion,

    pub(crate) engine: E,
    pub(crate) region_info: R,
}

/// The progress of a backup task
pub struct Progress<R: RegionInfoProvider> {
    store_id: u64,
    ranges: Vec<(Option<Key>, Option<Key>)>,
    next_index: usize,
    next_start: Option<Key>,
    end_key: Option<Key>,
    region_info: R,
    finished: bool,
}

impl<R: RegionInfoProvider> Progress<R> {
    fn new_with_range(
        store_id: u64,
        next_start: Option<Key>,
        end_key: Option<Key>,
        region_info: R,
    ) -> Self {
        let ranges = vec![(next_start, end_key)];
        Self::new_with_ranges(store_id, ranges, region_info)
    }

    fn new_with_ranges(
        store_id: u64,
        ranges: Vec<(Option<Key>, Option<Key>)>,
        region_info: R,
    ) -> Self {
        let mut prs = Progress {
            store_id,
            ranges,
            next_index: 0,
            next_start: None,
            end_key: None,
            region_info,
            finished: false,
        };
        prs.try_next();
        prs
    }

    /// try the next range. If all the ranges are consumed,
    /// set self.finish true.
    fn try_next(&mut self) {
        if self.ranges.len() > self.next_index {
            (self.next_start, self.end_key) = self.ranges[self.next_index].clone();

            self.next_index += 1;
        } else {
            self.finished = true;
        }
    }

    /// Forward the progress by `ranges` BackupRanges
    ///
    /// The size of the returned BackupRanges should <= `ranges`
    ///
    /// Notice: Returning an empty BackupRanges means that no leader region
    /// corresponding to the current range is sought. The caller should
    /// call `forward` again to seek regions for the next range.
    fn forward(&mut self, limit: usize) -> Option<Vec<BackupRange>> {
        if self.finished {
            return None;
        }
        let store_id = self.store_id;
        let (tx, rx) = mpsc::channel();
        let start_key_ = self
            .next_start
            .clone()
            .map_or_else(Vec::new, |k| k.into_encoded());

        let start_key = self.next_start.clone();
        let end_key = self.end_key.clone();
        let res = self.region_info.seek_region(
            &start_key_,
            Box::new(move |iter| {
                let mut count = 0;
                for info in iter {
                    let region = &info.region;
                    if let Some(end_key) = end_key.as_ref() {
                        let end_slice = end_key.as_encoded().as_slice();
                        if end_slice <= region.get_start_key() {
                            // We have reached the end.
                            // The range is defined as [start, end) so break if
                            // region start key is greater or equal to end key.
                            break;
                        }
                    }
                    if info.is_leader() {
                        let ekey = get_min_end_key(end_key.as_ref(), region);
                        let skey = get_max_start_key(start_key.as_ref(), region);
                        assert!(!(skey == ekey && ekey.is_some()), "{:?} {:?}", skey, ekey);
                        let leader = find_peer(region, store_id).unwrap().to_owned();
                        let backup_range = BackupRange {
                            start_key: skey,
                            end_key: ekey,
                            region: region.clone(),
                            leader,
                        };
                        tx.send(backup_range).unwrap();
                        count += 1;
                        if count >= limit {
                            break;
                        }
                    }
                }
            }),
        );
        if let Err(e) = res {
            // TODO: handle error.
            error!(?e; "backup seek region failed");
        }

        let branges: Vec<_> = rx.iter().collect();
        if let Some(b) = branges.last() {
            // The region's end key is empty means it is the last
            // region, we need to set the `finished` flag here in case
            // we run with `next_start` set to None
            if b.region.get_end_key().is_empty() || b.end_key == self.end_key {
                self.try_next();
            } else {
                self.next_start = b.end_key.clone();
            }
        } else {
            self.try_next();
        }
        Some(branges)
    }
}

impl<E: Engine, R: RegionInfoProvider + Clone + 'static> Endpoint<E, R> {
    pub fn new(
        store_id: u64,
        engine: E,
        region_info: R,
        config: BackupConfig,
        concurrency_manager: ConcurrencyManager,
        api_version: ApiVersion,
    ) -> Endpoint<E, R> {
        let pool = ControlThreadPool::new();
        let rt = utils::create_tokio_runtime(config.io_thread_size, "backup-io").unwrap();
        let config_manager = ConfigManager(Arc::new(RwLock::new(config)));
        let softlimit = SoftLimitKeeper::new(config_manager.clone());
        rt.spawn(softlimit.clone().run());
        Endpoint {
            store_id,
            engine,
            region_info,
            pool: RefCell::new(pool),
            io_pool: rt,
            softlimit,
            config_manager,
            concurrency_manager,
            api_version,
        }
    }

    pub fn get_config_manager(&self) -> ConfigManager {
        self.config_manager.clone()
    }

    fn get_config(&self) -> BackendConfig {
        BackendConfig {
            s3_multi_part_size: self.config_manager.0.read().unwrap().s3_multi_part_size.0 as usize,
            hdfs_config: HdfsConfig {
                hadoop_home: self.config_manager.0.read().unwrap().hadoop.home.clone(),
                linux_user: self
                    .config_manager
                    .0
                    .read()
                    .unwrap()
                    .hadoop
                    .linux_user
                    .clone(),
            },
        }
    }

    fn spawn_backup_worker(
        &self,
        prs: Arc<Mutex<Progress<R>>>,
        request: Request,
        saver_tx: async_channel::Sender<InMemBackupFiles>,
        resp_tx: UnboundedSender<BackupResponse>,
        _backend: Arc<dyn ExternalStorage>,
    ) {
        let start_ts = request.start_ts;
        let backup_ts = request.end_ts;
        let engine = self.engine.clone();
        let store_id = self.store_id;
        let concurrency_manager = self.concurrency_manager.clone();
        let batch_size = self.config_manager.0.read().unwrap().batch_size;
        let sst_max_size = self.config_manager.0.read().unwrap().sst_max_size.0;
        let limit = self.softlimit.limit();

        self.pool.borrow_mut().spawn(tikv_util::init_task_local(async move {
            loop {
                // when get the guard, release it until we finish scanning a batch,
                // because if we were suspended during scanning,
                // the region info have higher possibility to change (then we must compensate that by the fine-grained backup).
                let guard = limit.guard().await;
                if let Err(e) = guard {
                    warn!("failed to retrieve limit guard, omitting."; "err" => %e);
                };
                let batch = {
                    // Release lock as soon as possible.
                    // It is critical to speed up backup, otherwise workers are
                    // blocked by each other.
                    //
                    // If we use [tokio::sync::Mutex] here, until we give back the control flow to the scheduler
                    // or tasks waiting for the lock won't be waked up due to the characteristic of the runtime...
                    //
                    // The worst case is when using `noop` backend:
                    // the task seems never yielding and in fact the backup would executing sequentially.
                    //
                    // Anyway, even tokio itself doesn't recommend to use it unless the lock guard needs to be `Send`.
                    // (See https://tokio.rs/tokio/tutorial/shared-state)
                    // Use &mut and mark the type for making rust-analyzer happy.
                    let progress: &mut Progress<_> = &mut prs.lock().unwrap();
                    match progress.forward(batch_size) {
                        Some(batch) => batch,
                        None => return,
                    }
                };

                for brange in batch {
                    // wake up the scheduler for each loop for awaking tasks waiting for some lock or channels.
                    // because the softlimit permit is held by current task, there isn't risk of being suspended for long time.
                    tokio::task::yield_now().await;
                    let engine = engine.clone();
                    if request.cancel.load(Ordering::SeqCst) {
                        warn!("backup task has canceled"; "range" => ?brange);
                        return;
                    }
                    let ct = to_sst_compression_type(request.compression_type);

                    let writer_builder = BackupWriterBuilder::new(
                        store_id,
                        request.limiter.clone(),
                        brange.region.clone(),
                        ct,
                        request.compression_level,
                        sst_max_size,
                        request.cipher.clone(),
                    );
                    let stat = brange
                        .backup(
                            writer_builder,
                            engine,
                            concurrency_manager.clone(),
                            backup_ts,
                            start_ts,
                            saver_tx.clone(),
                            _backend.name(),
                        )
                        .await;
                    match stat {
                        Err(err) => {
                            error_unknown!(%err; "error during backup"; "region" => ?brange.region,);
                            let mut resp = BackupResponse::new();
                            resp.set_error(err.into());
                            if let Err(err) =  resp_tx.unbounded_send(resp) {
                                warn!("failed to send response"; "err" => ?err)
                            }
                        }
                        Ok(stat) => {
                            BACKUP_RAW_EXPIRED_COUNT.inc_by(stat.data.raw_value_tombstone as u64);
                            // TODO: maybe add the stat to metrics?
                            debug!("backup region finish";
                            "region" => ?brange.region,
                            "details" => ?stat);
                        }
                    }
                }
            }
        }));
    }

    fn get_progress_by_req(
        &self,
        request: &Request,
        codec: KeyValueCodec,
    ) -> Arc<Mutex<Progress<R>>> {
        if request.sub_ranges.is_empty() {
            let start_key = codec.encode_backup_key(request.start_key.clone());
            let end_key = codec.encode_backup_key(request.end_key.clone());
            Arc::new(Mutex::new(Progress::new_with_range(
                self.store_id,
                start_key,
                end_key,
                self.region_info.clone(),
            )))
        } else {
            let mut ranges = Vec::with_capacity(request.sub_ranges.len());
            for k in &request.sub_ranges {
                let start_key = codec.encode_backup_key(k.start_key.clone());
                let end_key = codec.encode_backup_key(k.end_key.clone());
                ranges.push((start_key, end_key));
            }
            Arc::new(Mutex::new(Progress::new_with_ranges(
                self.store_id,
                ranges,
                self.region_info.clone(),
            )))
        }
    }

    pub fn handle_backup_task(&self, task: Task) {
        let Task { request, resp } = task;
        let codec = KeyValueCodec::new(self.api_version, request.dst_api_ver);

        let prs = self.get_progress_by_req(&request, codec);

        let backend = match create_storage(&request.backend, self.get_config()) {
            Ok(backend) => backend,
            Err(err) => {
                error_unknown!(?err; "backup create storage failed");
                let mut response = BackupResponse::default();
                response.set_error(crate::Error::Io(err).into());
                if let Err(err) = resp.unbounded_send(response) {
                    error_unknown!(?err; "backup failed to send response");
                }
                return;
            }
        };
        let backend = Arc::<dyn ExternalStorage>::from(backend);
        let concurrency = self.config_manager.0.read().unwrap().num_threads;
        self.pool.borrow_mut().adjust_with(concurrency);
        // make the buffer small enough to implement back pressure.
        let (tx, rx) = async_channel::bounded(1);
        for _ in 0..concurrency {
            self.spawn_backup_worker(
                prs.clone(),
                request.clone(),
                tx.clone(),
                resp.clone(),
                backend.clone(),
            );
            self.io_pool.spawn(save_backup_file_worker(
                rx.clone(),
                resp.clone(),
                backend.clone(),
                codec,
            ));
        }
    }
}

impl<E: Engine, R: RegionInfoProvider + Clone + 'static> Runnable for Endpoint<E, R> {
    type Task = Task;

    fn run(&mut self, task: Task) {
        if task.has_canceled() {
            warn!("backup task has canceled"; "task" => %task);
            return;
        }
        info!("run backup task"; "task" => %task);
        self.handle_backup_task(task);
    }
}

/// Get the min end key from the given `end_key` and `Region`'s end key.
fn get_min_end_key(end_key: Option<&Key>, region: &Region) -> Option<Key> {
    let region_end = if region.get_end_key().is_empty() {
        None
    } else {
        Some(Key::from_encoded_slice(region.get_end_key()))
    };
    if region.get_end_key().is_empty() {
        end_key.cloned()
    } else if end_key.is_none() {
        region_end
    } else {
        let end_slice = end_key.as_ref().unwrap().as_encoded().as_slice();
        if end_slice < region.get_end_key() {
            end_key.cloned()
        } else {
            region_end
        }
    }
}

/// Get the max start key from the given `start_key` and `Region`'s start key.
fn get_max_start_key(start_key: Option<&Key>, region: &Region) -> Option<Key> {
    let region_start = if region.get_start_key().is_empty() {
        None
    } else {
        Some(Key::from_encoded_slice(region.get_start_key()))
    };
    if let Some(start_key) = start_key {
        let start_slice = start_key.as_encoded().as_slice();
        if start_slice >= region.get_start_key() {
            return Some(start_key.clone());
        }
    }
    region_start
}

/// Construct an backup file name based on the given store id, region, range
/// start key and local unix timestamp. A name consists with five parts: store
/// id, region_id, a epoch version, the hash of range start key and timestamp.
/// range start key is used to keep the unique file name for file, to handle
/// different tables exists on the same region. local unix timestamp is used to
/// keep the unique file name for file, to handle receive the same request after
/// connection reset.
pub fn backup_file_name(
    store_id: u64,
    region: &Region,
    key: Option<String>,
    storage_name: &str,
) -> String {
    let start = SystemTime::now();
    let since_the_epoch = start
        .duration_since(UNIX_EPOCH)
        .expect("Time went backwards");

    match (key, storage_name) {
        // See https://github.com/pingcap/tidb/issues/30087
        // To avoid 503 Slow Down error, if the backup storage is s3,
        // organize the backup files by store_id (use slash (/) as delimiter).
        (Some(k), aws::STORAGE_NAME | external_storage::local::STORAGE_NAME) => {
            format!(
                "{}/{}_{}_{}_{}",
                store_id,
                region.get_id(),
                region.get_region_epoch().get_version(),
                k,
                since_the_epoch.as_millis()
            )
        }
        (Some(k), _) => {
            format!(
                "{}_{}_{}_{}_{}",
                store_id,
                region.get_id(),
                region.get_region_epoch().get_version(),
                k,
                since_the_epoch.as_millis()
            )
        }

        (None, aws::STORAGE_NAME | external_storage::local::STORAGE_NAME) => {
            format!(
                "{}/{}_{}",
                store_id,
                region.get_id(),
                region.get_region_epoch().get_version()
            )
        }
        (None, _) => {
            format!(
                "{}_{}_{}",
                store_id,
                region.get_id(),
                region.get_region_epoch().get_version()
            )
        }
    }
}

// convert BackupCompresionType to rocks db DBCompressionType
fn to_sst_compression_type(ct: CompressionType) -> Option<SstCompressionType> {
    match ct {
        CompressionType::Lz4 => Some(SstCompressionType::Lz4),
        CompressionType::Snappy => Some(SstCompressionType::Snappy),
        CompressionType::Zstd => Some(SstCompressionType::Zstd),
        CompressionType::Unknown => None,
    }
}

/// warp a Option<Key> to the redact wrapper.
fn redact_option_key(key: &Option<Key>) -> log_wrappers::Value<'_> {
    match key {
        None => log_wrappers::Value::key(b""),
        Some(key) => log_wrappers::Value::key(key.as_encoded().as_slice()),
    }
}

#[cfg(test)]
pub mod tests {
    use std::{sync::Mutex, time::Duration};

    use file_system::IoRateLimiter;
    use kvproto::metapb;
    use raft::StateRole;
    use raftstore::coprocessor::{RegionCollector, Result as CopResult, SeekRegionCallback};
    use tempfile::TempDir;
    use tikv::storage::{BTreeEngine, TestEngineBuilder};
    use tikv_util::{config::ReadableSize, store::new_peer};
    use tokio::time;

    use super::*;

    #[derive(Clone)]
    pub struct MockRegionInfoProvider {
        regions: Arc<Mutex<RegionCollector>>,
        cancel: Option<Arc<AtomicBool>>,
        need_encode_key: bool,
    }

    impl MockRegionInfoProvider {
        pub fn new(encode_key: bool) -> Self {
            MockRegionInfoProvider {
                regions: Arc::new(Mutex::new(RegionCollector::default())),
                cancel: None,
                need_encode_key: encode_key,
            }
        }
        pub fn set_regions(&self, regions: Vec<(Vec<u8>, Vec<u8>, u64)>) {
            let mut map = self.regions.lock().unwrap();
            for (mut start_key, mut end_key, id) in regions {
                if !start_key.is_empty() {
                    if self.need_encode_key {
                        start_key = Key::from_raw(&start_key).into_encoded();
                    } else {
                        start_key = Key::from_encoded(start_key).into_encoded();
                    }
                }
                if !end_key.is_empty() {
                    if self.need_encode_key {
                        end_key = Key::from_raw(&end_key).into_encoded();
                    } else {
                        end_key = Key::from_encoded(end_key).into_encoded();
                    }
                }
                let mut r = metapb::Region::default();
                r.set_id(id);
                r.set_start_key(start_key.clone());
                r.set_end_key(end_key);
                r.mut_peers().push(new_peer(1, 1));
                map.create_region(r, StateRole::Leader);
            }
        }
    }

    impl RegionInfoProvider for MockRegionInfoProvider {
        fn seek_region(&self, from: &[u8], callback: SeekRegionCallback) -> CopResult<()> {
            let from = from.to_vec();
            let regions = self.regions.lock().unwrap();
            if let Some(c) = self.cancel.as_ref() {
                c.store(true, Ordering::SeqCst);
            }
            regions.handle_seek_region(from, callback);
            Ok(())
        }
    }

    pub fn new_endpoint() -> (TempDir, Endpoint<BTreeEngine, MockRegionInfoProvider>) {
        new_endpoint_with_limiter(None, ApiVersion::V1)
    }

    pub fn new_endpoint_with_limiter(
        limiter: Option<Arc<IoRateLimiter>>,
        api_version: ApiVersion,
    ) -> (TempDir, Endpoint<BTreeEngine, MockRegionInfoProvider>) {
        let temp = TempDir::new().unwrap();
        let engine = TestEngineBuilder::new()
            .path(temp.path())
            .cfs([
                engine_traits::CF_DEFAULT,
                engine_traits::CF_LOCK,
                engine_traits::CF_WRITE,
            ])
            .io_rate_limiter(limiter)
            .api_version(api_version)
            .build()
            .unwrap();
        let concurrency_manager = ConcurrencyManager::new(1.into());
        (
            temp,
            Endpoint::new(
                1,
                engine,
                MockRegionInfoProvider::new(true),
                BackupConfig {
                    num_threads: 4,
                    batch_size: 8,
                    sst_max_size: ReadableSize::mb(144),
                    ..Default::default()
                },
                concurrency_manager,
                api_version,
            ),
        )
    }

    #[test]
    fn test_control_thread_pool_adjust_keep_tasks() {
        use std::thread::sleep;

        let counter = Arc::new(AtomicU32::new(0));
        let mut pool = ControlThreadPool::new();
        pool.adjust_with(3);

        for i in 0..8 {
            let ctr = counter.clone();
            pool.spawn(async move {
                time::sleep(Duration::from_millis(100)).await;
                ctr.fetch_or(1 << i, Ordering::SeqCst);
            });
        }

        sleep(Duration::from_millis(150));
        pool.adjust_with(4);

        for i in 8..16 {
            let ctr = counter.clone();
            pool.spawn(async move {
                time::sleep(Duration::from_millis(100)).await;
                ctr.fetch_or(1 << i, Ordering::SeqCst);
            });
        }

        sleep(Duration::from_millis(250));
        assert_eq!(counter.load(Ordering::SeqCst), 0xffff);
    }

    #[test]
    fn test_s3_config() {
        let (_tmp, endpoint) = new_endpoint();
        assert_eq!(
            endpoint.config_manager.0.read().unwrap().s3_multi_part_size,
            ReadableSize::mb(5)
        );
    }

    #[test]
    fn test_seek_range() {
        let (_tmp, endpoint) = new_endpoint();

        endpoint.region_info.set_regions(vec![
            (b"".to_vec(), b"1".to_vec(), 1),
            (b"1".to_vec(), b"2".to_vec(), 2),
            (b"3".to_vec(), b"4".to_vec(), 3),
            (b"7".to_vec(), b"9".to_vec(), 4),
            (b"9".to_vec(), b"".to_vec(), 5),
        ]);
        // Test seek backup range.
        let test_seek_backup_range =
            |start_key: &[u8], end_key: &[u8], expect: Vec<(&[u8], &[u8])>| {
                let start_key = (!start_key.is_empty()).then_some(Key::from_raw(start_key));
                let end_key = (!end_key.is_empty()).then_some(Key::from_raw(end_key));
                let mut prs = Progress::new_with_range(
                    endpoint.store_id,
                    start_key,
                    end_key,
                    endpoint.region_info.clone(),
                );

                let mut ranges = Vec::with_capacity(expect.len());
                while ranges.len() != expect.len() {
                    let n = (rand::random::<usize>() % 3) + 1;
                    let mut r = prs.forward(n).unwrap();
                    // The returned backup ranges should <= n
                    assert!(r.len() <= n);

                    if r.is_empty() {
                        // if return a empty vec then the progress is finished
                        assert_eq!(
                            ranges.len(),
                            expect.len(),
                            "got {:?}, expect {:?}",
                            ranges,
                            expect
                        );
                    }
                    ranges.append(&mut r);
                }

                for (a, b) in ranges.into_iter().zip(expect) {
                    assert_eq!(
                        a.start_key.map_or_else(Vec::new, |k| k.into_raw().unwrap()),
                        b.0
                    );
                    assert_eq!(
                        a.end_key.map_or_else(Vec::new, |k| k.into_raw().unwrap()),
                        b.1
                    );
                }
            };

        // Backup range from case.0 to case.1,
        // the case.2 is the expected results.
        type Case<'a> = (&'a [u8], &'a [u8], Vec<(&'a [u8], &'a [u8])>);

        let case: Vec<Case<'_>> = vec![
            (b"", b"1", vec![(b"", b"1")]),
            (b"", b"2", vec![(b"", b"1"), (b"1", b"2")]),
            (b"1", b"2", vec![(b"1", b"2")]),
            (b"1", b"3", vec![(b"1", b"2")]),
            (b"1", b"4", vec![(b"1", b"2"), (b"3", b"4")]),
            (b"4", b"6", vec![]),
            (b"4", b"5", vec![]),
            (b"2", b"7", vec![(b"3", b"4")]),
            (b"7", b"8", vec![(b"7", b"8")]),
            (b"3", b"", vec![(b"3", b"4"), (b"7", b"9"), (b"9", b"")]),
            (b"5", b"", vec![(b"7", b"9"), (b"9", b"")]),
            (b"7", b"", vec![(b"7", b"9"), (b"9", b"")]),
            (b"8", b"91", vec![(b"8", b"9"), (b"9", b"91")]),
            (b"8", b"", vec![(b"8", b"9"), (b"9", b"")]),
            (
                b"",
                b"",
                vec![
                    (b"", b"1"),
                    (b"1", b"2"),
                    (b"3", b"4"),
                    (b"7", b"9"),
                    (b"9", b""),
                ],
            ),
        ];
        for (start_key, end_key, ranges) in case {
            test_seek_backup_range(start_key, end_key, ranges);
        }
    }

    #[test]
    fn test_seek_ranges() {
        let (_tmp, endpoint) = new_endpoint();

        endpoint.region_info.set_regions(vec![
            (b"".to_vec(), b"1".to_vec(), 1),
            (b"1".to_vec(), b"2".to_vec(), 2),
            (b"3".to_vec(), b"4".to_vec(), 3),
            (b"7".to_vec(), b"9".to_vec(), 4),
            (b"9".to_vec(), b"".to_vec(), 5),
        ]);
        // Test seek backup range.
        let test_seek_backup_ranges =
            |sub_ranges: Vec<(&[u8], &[u8])>, expect: Vec<(&[u8], &[u8])>| {
                let mut ranges = Vec::with_capacity(sub_ranges.len());
                for &(start_key, end_key) in &sub_ranges {
                    let start_key = (!start_key.is_empty()).then_some(Key::from_raw(start_key));
                    let end_key = (!end_key.is_empty()).then_some(Key::from_raw(end_key));
                    ranges.push((start_key, end_key));
                }
                let mut prs = Progress::new_with_ranges(
                    endpoint.store_id,
                    ranges,
                    endpoint.region_info.clone(),
                );

                let mut ranges = Vec::with_capacity(expect.len());
                loop {
                    let n = (rand::random::<usize>() % 3) + 1;
                    let mut r = match prs.forward(n) {
                        None => break,
                        Some(r) => r,
                    };
                    // The returned backup ranges should <= n
                    assert!(r.len() <= n);

                    if !r.is_empty() {
                        ranges.append(&mut r);
                    }
                }

                for (a, b) in ranges.into_iter().zip(expect) {
                    assert_eq!(
                        a.start_key.map_or_else(Vec::new, |k| k.into_raw().unwrap()),
                        b.0
                    );
                    assert_eq!(
                        a.end_key.map_or_else(Vec::new, |k| k.into_raw().unwrap()),
                        b.1
                    );
                }
            };

        // Backup range from case.0 to case.1,
        // the case.2 is the expected results.
        type Case<'a> = (Vec<(&'a [u8], &'a [u8])>, Vec<(&'a [u8], &'a [u8])>);

        let case: Vec<Case<'_>> = vec![
            (
                vec![(b"", b"1"), (b"1", b"2")],
                vec![(b"", b"1"), (b"1", b"2")],
            ),
            (
                vec![(b"", b"2"), (b"3", b"4")],
                vec![(b"", b"1"), (b"1", b"2"), (b"3", b"4")],
            ),
            (
                vec![(b"7", b"8"), (b"8", b"9")],
                vec![(b"7", b"8"), (b"8", b"9")],
            ),
            (
                vec![(b"8", b"9"), (b"6", b"8")],
                vec![(b"8", b"9"), (b"7", b"8")],
            ),
            (
                vec![(b"8", b"85"), (b"88", b"89"), (b"7", b"8")],
                vec![(b"8", b"85"), (b"88", b"89"), (b"7", b"8")],
            ),
            (
                vec![(b"8", b"85"), (b"", b"35"), (b"88", b"89"), (b"7", b"8")],
                vec![
                    (b"8", b"85"),
                    (b"", b"1"),
                    (b"1", b"2"),
                    (b"3", b"35"),
                    (b"88", b"89"),
                    (b"7", b"8"),
                ],
            ),
            (vec![(b"", b"1")], vec![(b"", b"1")]),
            (vec![(b"", b"2")], vec![(b"", b"1"), (b"1", b"2")]),
            (vec![(b"1", b"2")], vec![(b"1", b"2")]),
            (vec![(b"1", b"3")], vec![(b"1", b"2")]),
            (vec![(b"1", b"4")], vec![(b"1", b"2"), (b"3", b"4")]),
            (vec![(b"4", b"5")], vec![]),
            (vec![(b"4", b"6")], vec![]),
            (vec![(b"4", b"6"), (b"6", b"7")], vec![]),
            (vec![(b"2", b"3"), (b"4", b"6"), (b"6", b"7")], vec![]),
            (vec![(b"2", b"7")], vec![(b"3", b"4")]),
            (vec![(b"7", b"8")], vec![(b"7", b"8")]),
            (
                vec![(b"3", b"")],
                vec![(b"3", b"4"), (b"7", b"9"), (b"9", b"")],
            ),
            (vec![(b"5", b"")], vec![(b"7", b"9"), (b"9", b"")]),
            (vec![(b"7", b"")], vec![(b"7", b"9"), (b"9", b"")]),
            (vec![(b"8", b"91")], vec![(b"8", b"9"), (b"9", b"91")]),
            (vec![(b"8", b"")], vec![(b"8", b"9"), (b"9", b"")]),
            (
                vec![(b"", b"")],
                vec![
                    (b"", b"1"),
                    (b"1", b"2"),
                    (b"3", b"4"),
                    (b"7", b"9"),
                    (b"9", b""),
                ],
            ),
        ];
        for (ranges, expect_ranges) in case {
            test_seek_backup_ranges(ranges, expect_ranges);
        }
    }

    fn fake_empty_marker() -> Vec<super::BackupRange> {
        vec![super::BackupRange {
            start_key: None,
            end_key: None,
            region: Region::new(),
            leader: Peer::new(),
        }]
    }

    #[test]
    fn test_seek_ranges_2() {
        let (_tmp, endpoint) = new_endpoint();

        endpoint.region_info.set_regions(vec![
            (b"2".to_vec(), b"4".to_vec(), 1),
            (b"6".to_vec(), b"8".to_vec(), 2),
        ]);
        let sub_ranges: Vec<(&[u8], &[u8])> = vec![(b"1", b"11"), (b"3", b"7"), (b"8", b"9")];
        let expect: Vec<(&[u8], &[u8])> = vec![(b"", b""), (b"3", b"4"), (b"6", b"7"), (b"", b"")];

        let mut ranges = Vec::with_capacity(sub_ranges.len());
        for &(start_key, end_key) in &sub_ranges {
            let start_key = (!start_key.is_empty()).then_some(Key::from_raw(start_key));
            let end_key = (!end_key.is_empty()).then_some(Key::from_raw(end_key));
            ranges.push((start_key, end_key));
        }
        let mut prs =
            Progress::new_with_ranges(endpoint.store_id, ranges, endpoint.region_info.clone());

        let mut ranges = Vec::with_capacity(expect.len());
        loop {
            let n = (rand::random::<usize>() % 2) + 1;
            let mut r = match prs.forward(n) {
                None => break,
                Some(r) => r,
            };
            // The returned backup ranges should <= n
            assert!(r.len() <= n);

            if !r.is_empty() {
                ranges.append(&mut r);
            } else {
                // append the empty marker
                ranges.append(&mut fake_empty_marker());
            }
        }

        assert!(ranges.len() == expect.len());
        for (a, b) in ranges.into_iter().zip(expect) {
            assert_eq!(
                a.start_key.map_or_else(Vec::new, |k| k.into_raw().unwrap()),
                b.0
            );
            assert_eq!(
                a.end_key.map_or_else(Vec::new, |k| k.into_raw().unwrap()),
                b.1
            );
        }
    }

    #[test]
    fn test_backup_file_name() {
        let region = metapb::Region::default();
        let store_id = 1;
        let test_cases = ["s3", "local", "gcs", "azure", "hdfs"];
        let test_target = [
            "1/0_0_000",
            "1/0_0_000",
            "1_0_0_000",
            "1_0_0_000",
            "1_0_0_000",
        ];

        let delimiter = "_";
        for (storage_name, target) in test_cases.iter().zip(test_target.iter()) {
            let key = Some(String::from("000"));
            let filename = backup_file_name(store_id, &region, key, storage_name);

            let mut prefix_arr: Vec<&str> = filename.split(delimiter).collect();
            prefix_arr.remove(prefix_arr.len() - 1);

            assert_eq!(target.to_string(), prefix_arr.join(delimiter));
        }

        let test_target = ["1/0_0", "1/0_0", "1_0_0", "1_0_0", "1_0_0"];
        for (storage_name, target) in test_cases.iter().zip(test_target.iter()) {
            let key = None;
            let filename = backup_file_name(store_id, &region, key, storage_name);
            assert_eq!(target.to_string(), filename);
        }
    }
}