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

#[cfg(test)]
use std::io::Write;
#[cfg(any(test, feature = "testexport"))]
use std::sync::atomic::AtomicBool;
use std::{
    borrow::Cow,
    collections::{HashMap, VecDeque},
    fmt, fs,
    fs::File,
    io::{BufReader, Read, Seek, SeekFrom},
    path::{Path, PathBuf},
    sync::{
        Arc, RwLock,
        atomic::{AtomicU32, Ordering},
    },
    thread,
    thread::JoinHandle,
    time::Duration,
};

use arc_swap::ArcSwap;
use byteorder::{ByteOrder, LittleEndian};
use bytes::{BufMut, BytesMut};
use file_system::IoRateLimiter;
use kvengine::dfs::{Dfs, MultiPartUploader};
use protobuf::Message;
use rfenginepb::{
    KeySpaceBackupMeta, RaftLogBackupFile, RaftLogFile, StoreBackupMeta, StoreRaftLogBackupMeta,
};
use slog_global::*;
use tikv_util::{
    backoff,
    errors::{Context as _, IoError},
    mpsc::{Receiver, SendError, Sender},
    time::Instant,
};

macro_rules! try_force_stop {
    ($worker:expr) => {{
        #[cfg(any(test, feature = "testexport"))]
        if $worker.force_stop.load(Ordering::SeqCst) {
            info!("compact worker force stop");
            return;
        }
    }};
}

macro_rules! try_force_stop_err {
    ($worker:expr) => {{
        #[cfg(any(test, feature = "testexport"))]
        if $worker.force_stop.load(Ordering::SeqCst) {
            info!("compact worker force stop");
            return Err(Error::CompactWorkerForceStop);
        }
    }};
}

use crate::{
    compact::SubCompactionPool,
    manifest::{Manifest, PeerFile},
    metrics::{RFENGINE_BACKUP_COUNTER, RFENGINE_BACKUP_DURATION_HISTOGRAM},
    write_batch::PeerBatch,
    *,
};

const COMPACT_RETRY_TIMES: usize = 30;
#[cfg(test)]
const WRITE_RLOG_BATCH_SIZE: usize = 256 * 1024;
const SNAPSHOT_RLOG_META_RESERVED_SIZE: usize = 10 * 1024 * 1024;

pub(crate) struct WorkerHandle {
    pub(crate) task_sender: Sender<CompactTask>,
    pub(crate) handle: Option<JoinHandle<()>>,
}

impl WorkerHandle {
    pub(crate) fn try_send(&self, task: CompactTask) {
        if let Err(SendError(_)) = self.task_sender.send(task) {
            warn!("failed to send task");
        }
    }
}

struct WeightedPeerBatch {
    peer_batch: PeerBatch,
    estimated_size: usize,
}

enum SnapshotRlogObjectBuilder {
    Inline {
        key: String,
        object: BytesMut,
        max_object_size: usize,
    },
    Multipart {
        uploader: MultiPartUploader,
        max_object_size: usize,
    },
}

impl SnapshotRlogObjectBuilder {
    fn new(
        dfs: Option<&Arc<dyn Dfs>>,
        key: String,
        estimated_rlog_size: usize,
        multipart_threshold: usize,
        multipart_part_size: usize,
        multipart_concurrency: usize,
        max_object_size: usize,
    ) -> Result<Self> {
        let estimated_object_size =
            estimated_rlog_size.saturating_add(SNAPSHOT_RLOG_META_RESERVED_SIZE);
        if should_use_snapshot_rlog_multipart_upload(estimated_object_size, multipart_threshold) {
            let dfs = dfs.ok_or_else(|| {
                Error::Dfs("snapshot rlog multipart upload requires DFS".to_string())
            })?;
            return Ok(Self::Multipart {
                uploader: MultiPartUploader::new(
                    Arc::clone(dfs),
                    key,
                    multipart_part_size,
                    multipart_concurrency,
                )
                .map_err(Error::Dfs)?,
                max_object_size,
            });
        }

        Ok(Self::Inline {
            key,
            object: BytesMut::with_capacity(estimated_object_size.min(max_object_size)),
            max_object_size,
        })
    }

    fn len(&self) -> usize {
        match self {
            SnapshotRlogObjectBuilder::Inline { object, .. } => object.len(),
            SnapshotRlogObjectBuilder::Multipart { uploader, .. } => uploader.len(),
        }
    }

    fn max_object_size(&self) -> usize {
        match self {
            SnapshotRlogObjectBuilder::Inline {
                max_object_size, ..
            }
            | SnapshotRlogObjectBuilder::Multipart {
                max_object_size, ..
            } => *max_object_size,
        }
    }

    fn put_slice(&mut self, data: &[u8]) -> Result<()> {
        let object_size = self
            .len()
            .checked_add(data.len())
            .ok_or(Error::SnapshotOversize(u64::MAX))?;
        if object_size > self.max_object_size() {
            return Err(Error::SnapshotOversize(object_size as u64));
        }

        match self {
            SnapshotRlogObjectBuilder::Inline { object, .. } => {
                object.put_slice(data);
                Ok(())
            }
            SnapshotRlogObjectBuilder::Multipart { uploader, .. } => {
                uploader.upload_part(data).map_err(Error::Dfs)
            }
        }
    }

    fn finish(self) -> Result<PreparedSnapshotRlogObject> {
        match self {
            SnapshotRlogObjectBuilder::Inline { key, object, .. } => {
                Ok(PreparedSnapshotRlogObject::Inline {
                    key,
                    data: object.freeze(),
                })
            }
            SnapshotRlogObjectBuilder::Multipart { uploader, .. } => Ok(
                PreparedSnapshotRlogObject::Multipart(uploader.finish().map_err(Error::Dfs)?),
            ),
        }
    }
}

fn build_sub_batches_for_compaction(
    peer_batches: Vec<PeerBatch>,
    compaction_batch_size: usize,
) -> Vec<Vec<PeerBatch>> {
    assert!(
        compaction_batch_size > 0,
        "compaction_batch_size must be non-zero"
    );
    if peer_batches.is_empty() {
        return vec![];
    }

    let peer_batch_count = peer_batches.len();
    let sub_batch_count = peer_batch_count.div_ceil(compaction_batch_size);

    // Keep the old compact-worker behavior: even if a decoded/cached peer batch
    // still carries raft logs already covered by `truncated_idx`, trim them
    // before we estimate work or write new `.rlog` files. Compute the encoded
    // size in the same pass so weighted planning does not need to walk the peer
    // batches repeatedly.
    let mut weighted_peer_batches = Vec::with_capacity(peer_batch_count);
    for mut peer_batch in peer_batches {
        peer_batch.truncate(peer_batch.truncated_idx);
        weighted_peer_batches.push(WeightedPeerBatch {
            estimated_size: peer_batch.encoded_len(),
            peer_batch,
        });
    }

    if sub_batch_count <= 1 {
        return vec![
            weighted_peer_batches
                .into_iter()
                .map(|weighted| weighted.peer_batch)
                .collect(),
        ];
    }

    // Keep the configured peer-count cap, but balance the actual work by
    // encoded bytes so a few large peers are less likely to dominate the tail
    // of one epoch.
    weighted_peer_batches.sort_unstable_by(|lhs, rhs| {
        rhs.estimated_size
            .cmp(&lhs.estimated_size)
            .then_with(|| lhs.peer_batch.peer_id.cmp(&rhs.peer_batch.peer_id))
    });

    // Distribute the largest peer batches first in round-robin order. This is
    // simpler to review than maintaining a heap, while still spreading large
    // peers across different sub-batches. Because `sub_batch_count` is derived
    // from `ceil(total_peers / compaction_batch_size)`, round-robin assignment
    // also keeps every sub-batch within the configured peer-count cap.
    let mut sub_batches = (0..sub_batch_count).map(|_| Vec::new()).collect::<Vec<_>>();
    for (idx, weighted_peer_batch) in weighted_peer_batches.into_iter().enumerate() {
        sub_batches[idx % sub_batch_count].push(weighted_peer_batch.peer_batch);
    }

    sub_batches
        .into_iter()
        .filter(|batch| !batch.is_empty())
        .collect()
}

pub(crate) struct CompactWorker {
    dir: PathBuf,
    manifest: Manifest,
    task_rx: Receiver<CompactTask>,
    snap_tx: Sender<ObjectStorageTask>,
    #[cfg(test)]
    buf: Vec<u8>,
    compacted_epoch: Arc<AtomicU32>,
    epoch_rotate_len: usize,
    delay_compaction_epoches: u32,

    // Used to cache small rlogs to reduce disk IO when taking snapshot.
    rlog_cache: Arc<RwLock<RlogCache>>,
    rlog_compression_type: CompressionType,
    snapshot_rlog_multipart_upload_threshold: usize,
    snapshot_rlog_multipart_upload_part_size: usize,
    snapshot_rlog_multipart_upload_concurrency: usize,
    snapshot_rlog_max_object_size: usize,
    snapshot_rlog_multipart_dfs: Option<Arc<dyn Dfs>>,

    rate_limiter: Option<Arc<IoRateLimiter>>,

    // Size limit (bytes) for splitting rlog files during WAL
    // compaction.
    rlog_file_size: u32,

    #[cfg(any(test, feature = "testexport"))]
    force_stop: Arc<AtomicBool>,

    // Shared per-peer Rlog files; see `RfEngineCore::peer_rlog_files`.
    peer_rlog_files: Option<Arc<ArcSwap<HashMap<u64, VecDeque<PeerFile>>>>>,
    sub_compaction_pool: SubCompactionPool,
    compaction_batch_size: usize,

    // cache write batches used for compaction, each write belongs to exact 1 epoch.
    cached_write_batches: VecDeque<CachedWriteBatch>,
    // the write batch for the current epoch.
    current_batch: CachedWriteBatch,
}

#[derive(Default)]
struct CachedWriteBatch {
    epoch_id: u32,
    wb: WriteBatch,
}

impl CompactWorker {
    pub(crate) fn new(
        dir: PathBuf,
        cfg: &RfEngineConfig,
        task_rx: Receiver<CompactTask>,
        snap_tx: Sender<ObjectStorageTask>,
        manifest: Manifest,
        compacted_epoch: Arc<AtomicU32>,
        rlog_cache: Arc<RwLock<RlogCache>>,
        rlog_compression_type: CompressionType,
        snapshot_rlog_multipart_upload_threshold: usize,
        snapshot_rlog_multipart_upload_part_size: usize,
        snapshot_rlog_multipart_upload_concurrency: usize,
        snapshot_rlog_max_object_size: usize,
        snapshot_rlog_multipart_dfs: Option<Arc<dyn Dfs>>,
        rate_limiter: Option<Arc<IoRateLimiter>>,
        #[cfg(any(test, feature = "testexport"))] force_stop: Arc<AtomicBool>,
        peer_rlog_files: Arc<ArcSwap<HashMap<u64, VecDeque<PeerFile>>>>,
    ) -> Self {
        let epoch_rotate_len = cfg.epoch_rotate_len;
        let delay_compaction_epoches = cfg.delay_compaction_epoches as u32;
        let rlog_file_size = cfg.rlog_file_size.0 as u32;
        assert!(
            cfg.compaction_concurrency > 0,
            "compaction_concurrency must be non-zero"
        );
        assert!(
            cfg.compaction_batch_size > 0,
            "compaction_batch_size must be non-zero"
        );
        Self {
            dir,
            manifest,
            task_rx,
            snap_tx,
            #[cfg(test)]
            buf: vec![],
            compacted_epoch,
            epoch_rotate_len,
            delay_compaction_epoches,
            rlog_cache,
            rlog_compression_type,
            snapshot_rlog_multipart_upload_threshold,
            snapshot_rlog_multipart_upload_part_size,
            snapshot_rlog_multipart_upload_concurrency,
            snapshot_rlog_max_object_size,
            snapshot_rlog_multipart_dfs,
            rate_limiter,
            rlog_file_size,
            #[cfg(any(test, feature = "testexport"))]
            force_stop,
            peer_rlog_files: cfg.rlog_offloading_enabled().then_some(peer_rlog_files),
            sub_compaction_pool: SubCompactionPool::new(cfg.compaction_concurrency),
            compaction_batch_size: cfg.compaction_batch_size,
            cached_write_batches: VecDeque::new(),
            current_batch: CachedWriteBatch::default(),
        }
    }

    pub(crate) fn run(&mut self) {
        while let Ok(task) = self.task_rx.recv() {
            try_force_stop!(self);
            match task {
                CompactTask::Write { wb } => self.handle_write(&wb),
                CompactTask::Rotate {
                    epoch_id,
                    cache_wb_for_compact,
                } => self.handle_rotate(epoch_id, cache_wb_for_compact),
                CompactTask::PrepareSnapshot(request_token) => {
                    self.handle_prepare_snapshot_backup(request_token);
                }
                CompactTask::Close => return,
            }
        }
    }

    fn handle_compact_with_backoff(&mut self, epoch_id: u32, delayed_epoches: u32) {
        let mut back_off = backoff::ExponentialBackoff::new(
            Duration::from_secs(5),
            Duration::from_secs(60),
            COMPACT_RETRY_TIMES,
        );
        let engine_id = self.manifest.get_engine_id();
        loop {
            let next_delay = back_off.next_delay();
            if next_delay.is_err() {
                panic!("compact epoch {} retry times exceeded", epoch_id);
            }
            match self.compact(epoch_id, delayed_epoches) {
                Ok(_) => break,
                Err(Error::CompactWorkerForceStop) => return,
                Err(err) => {
                    error!(
                        "{}: failed to compact epoch {} retry_times: {}, {:?}",
                        engine_id,
                        epoch_id,
                        back_off.current_attempts(),
                        err
                    );
                    thread::sleep(next_delay.unwrap());
                }
            }
        }
    }

    fn handle_write(&mut self, batch: &[PeerBatch]) {
        // skip new write batch if the current write batch is not initialized.
        // This can happen at system start as we don't get full changes of the first
        // WAL.
        if self.current_batch.epoch_id == 0 {
            return;
        }
        for b in batch {
            // try to truncate trunctead_logs in cache write_batches.
            if b.truncated_idx > 0
                && b.truncated_idx
                    > self
                        .current_batch
                        .wb
                        .get_peer(b.peer_id)
                        .map_or(0, |p| p.truncated_idx)
            {
                for wb in &mut self.cached_write_batches {
                    if let Some(peer_wb) = wb.wb.peers.get_mut(&b.peer_id) {
                        if peer_wb.truncated_idx >= b.truncated_idx {
                            break;
                        }
                        peer_wb.truncate(b.truncated_idx);
                    }
                }
            }
            self.current_batch.wb.merge_peer(b.clone());
        }
    }

    fn handle_rotate(&mut self, epoch_id: u32, cache_wb_for_compact: bool) {
        try_force_stop!(self);
        if self.current_batch.epoch_id != 0 {
            assert_eq!(self.current_batch.epoch_id, epoch_id);
            let batch = std::mem::take(&mut self.current_batch);
            self.cached_write_batches.push_back(batch);
        }
        if cache_wb_for_compact {
            assert!(self.current_batch.wb.is_empty());
            self.current_batch.epoch_id = epoch_id + 1;
        }
        let engine_id = self.manifest.get_engine_id();
        let delay_epoches = self.delay_compaction_epoches;
        let compact_epoch = epoch_id.saturating_sub(delay_epoches);
        if compact_epoch <= self.manifest.epoch_id {
            // Do not compact at smaller epoch.
            info!(
                "{}: skip duplicated epoch compaction at epoch {}",
                engine_id, compact_epoch
            );
            return;
        }
        info!(
            "{}: handle rotate {}, compact {}",
            engine_id, epoch_id, compact_epoch
        );
        if let Err(err) = self.compact(compact_epoch, delay_epoches) {
            if let Error::CompactWorkerForceStop = err {
                return;
            }
            error!(
                "{}: failed to compact epoch {} {:?}",
                engine_id, compact_epoch, err
            );
            self.handle_compact_with_backoff(compact_epoch, delay_epoches);
            info!("handle compact {} success after retry", compact_epoch);
        }
    }

    fn compact(&mut self, epoch_id: u32, delayed_epoches: u32) -> Result<()> {
        try_force_stop_err!(self);
        let timer = Instant::now_coarse();
        let cached_wb = match self.cached_write_batches.front().map(|wb| wb.epoch_id) {
            Some(cached_epoch_id) if cached_epoch_id <= epoch_id => {
                assert_eq!(cached_epoch_id, epoch_id);
                Some(self.cached_write_batches.pop_front().unwrap())
            }
            _ => None,
        };
        let mut batch = if let Some(wb) = cached_wb {
            // only verify in memory WriteBatch with disk file in test environment.
            #[cfg(test)]
            self.check_async_compact_wb(epoch_id, &wb.wb);
            wb.wb
        } else {
            let mut wb = self.load_wal(epoch_id)?;
            for peer_batch in wb.peers.values_mut() {
                // update the truncated_index with the newest truncated index.
                for wb in self.cached_write_batches.iter().rev() {
                    if let Some(pb) = wb.wb.get_peer(peer_batch.peer_id) {
                        if pb.truncated_idx > peer_batch.truncated_idx {
                            peer_batch.truncate(pb.truncated_idx);
                        }
                        break;
                    }
                }
            }
            wb
        };
        for peer_batch in batch.peers.values_mut() {
            if peer_batch.get_keyspace_id().is_none() {
                peer_batch.meta.keyspace_id = self
                    .manifest
                    .peers
                    .get(&peer_batch.peer_id)
                    .map(|peer| peer.keyspace_id)
                    .unwrap_or_default();
            }
        }
        let mut change_set = rfenginepb::ChangeSet::default();
        change_set.set_epoch_id(epoch_id);
        change_set.set_delayed_epoches(delayed_epoches);
        let peer_batches = batch.into_vector();
        let peer_count = peer_batches.len();
        let sub_batches =
            build_sub_batches_for_compaction(peer_batches, self.compaction_batch_size);
        let mut result = self.sub_compaction_pool.compact_epoch(
            self.dir.as_path(),
            self.rate_limiter.clone(),
            self.rlog_cache.clone(),
            self.rlog_file_size,
            epoch_id,
            sub_batches,
        )?;
        for meta in result.metas.drain(..) {
            change_set.mut_peers().push(meta);
        }

        fail::fail_point!("rfengine::compact_before_sync_dir");
        file_system::sync_dir(self.dir.as_path()).ctx("compact_sync_dir")?;

        try_force_stop_err!(self);
        self.manifest.handle_compaction(change_set)?;

        fail::fail_point!("rfengine::compact_after_manifest_persist");

        // Update the peer_rlog_files that is shared with rfengine core.
        if let Some(peer_rlog_files) = self.peer_rlog_files.as_ref() {
            peer_rlog_files.store(Arc::new(self.manifest.peer_rlog_files()))
        }

        let engine_id = self.manifest.get_engine_id();
        let duration = timer.saturating_elapsed();
        let pending_tasks = self.task_rx.len();
        let cache_size = self.rlog_cache.read().unwrap().cache_size();
        info!(
            "{}: compact wal {}", engine_id, epoch_id;
            "delayed_epoches" => delayed_epoches,
            "generated_files" => result.generated_files,
            "cached_files" => result.cached_files,
            "cache_size" => cache_size,
            "takes" => ?duration,
            "pending_tasks" => pending_tasks,
            "sum of peer" => peer_count,
        );
        ENGINE_COMPACT_WAL_DURATION_HISTOGRAM.observe(duration.as_secs_f64());
        self.compacted_epoch.store(epoch_id, Ordering::SeqCst);
        Ok(())
    }

    fn load_wal(&self, epoch_id: u32) -> Result<WriteBatch> {
        let mut batch = WriteBatch::default();
        let mut it = WalIterator::new(&self.dir, epoch_id, self.epoch_rotate_len)?;
        it.iterate_batch(|data, _| {
            iterate_peer_batch(data, |region_batch| {
                batch.merge_peer(region_batch);
            });
        })?;
        Ok(batch)
    }

    #[cfg(test)]
    fn check_async_compact_wb(&self, epoch_id: u32, wb: &WriteBatch) {
        let mut wb_disk = self.load_wal(epoch_id).unwrap();
        assert_eq!(wb.peers.len(), wb_disk.peers.len());
        for (id, peer_data) in &wb.peers {
            let Some(peer_disk) = wb_disk.peers.get_mut(id) else {
                panic!("[epoch {}] peer {} not found in disk", epoch_id, *id);
            };

            if peer_data.truncated_idx > peer_disk.truncated_idx {
                peer_disk.truncate(peer_data.truncated_idx);
            }

            if peer_disk.meta.keyspace_id == 0 {
                peer_disk.meta.keyspace_id = peer_data.meta.keyspace_id;
            }
            if peer_data.meta != peer_disk.meta {
                panic!(
                    "[epoch {}] peer {} meta mismatch: mem: {:?}, disk: {:?}",
                    epoch_id, *id, peer_data.meta, peer_disk.meta
                );
            }

            if peer_data.raft_logs.len() != peer_disk.raft_logs.len() {
                panic!(
                    "[epoch {}] peer {} raft logs length mismatch: mem: {}, disk: {}",
                    epoch_id,
                    *id,
                    peer_data.raft_logs.len(),
                    peer_disk.raft_logs.len(),
                );
            }

            for (log_mem, log_disk) in peer_data.raft_logs.iter().zip(&peer_disk.raft_logs) {
                if log_mem != log_disk {
                    panic!(
                        "[epoch {}] peer {} raft logs mismatch: mem: {:?}, disk: {:?}",
                        epoch_id, *id, log_mem, log_disk,
                    );
                }
            }
        }
    }

    // rfenginepb::RaftLogFile format:
    // RlogHeader + [endoffset] + [RaftLogOp + checksum]
    #[cfg(test)]
    fn write_raft_log_files(
        &mut self,
        peer_batch: PeerBatch,
        epoch_id: u32,
    ) -> Result<Vec<(rfenginepb::RaftLogFile, bool /* is_cached */)>> {
        let PeerBatch {
            peer_id, raft_logs, ..
        } = peer_batch;
        let mut res = Vec::new();
        let mut i = 0;
        // Split raft logs into multiple files if the raft log offset
        // exceeds size limit.
        while i < raft_logs.len() {
            let mut count = 0;
            let mut log_end_off: u32 = 0;
            while i + count < raft_logs.len() {
                let op = &raft_logs[i + count];
                let size = (op.encoded_len() as u32) + 4 /* checksum */;
                if let Some(new_off) = log_end_off.checked_add(size) {
                    if new_off <= self.rlog_file_size {
                        log_end_off = new_off;
                        count += 1;
                        continue;
                    }
                }
                break;
            }
            if count == 0 {
                // A legal Raft entry may exceed the target `rlog_file_size`.
                // Write it alone in an oversized rlog file.
                count = 1;
            }
            let (meta, is_cached) =
                self.write_raft_log_file(peer_id, &raft_logs, i, count, epoch_id)?;
            res.push((meta, is_cached));
            i += count;
        }
        Ok(res)
    }

    #[cfg(test)]
    fn write_raft_log_file(
        &mut self,
        peer_id: u64,
        raft_logs: &VecDeque<RaftLogOp>,
        start_idx: usize,
        count: usize,
        epoch_id: u32,
    ) -> Result<(rfenginepb::RaftLogFile, bool /* is_cached */)> {
        let first = raft_logs[start_idx].index;
        let last = raft_logs[start_idx + count - 1].index;
        let last_term = raft_logs[start_idx + count - 1].term;
        let filename = raft_log_file_name(self.dir.as_path(), peer_id, first, last);
        self.buf.truncate(0);
        let header = RlogHeader::new(count as u32);
        header.encode_to(&mut self.buf);
        // Write index first to make the file addressable.
        let mut log_end_off = 0;
        for idx in start_idx..(start_idx + count) {
            let rlog = &raft_logs[idx];
            log_end_off += rlog.encoded_len() as u32 + 4 /* checksum */;
            self.buf.put_u32_le(log_end_off);
        }
        for idx in start_idx..(start_idx + count) {
            let rlog = &raft_logs[idx];
            let origin_len = self.buf.len();
            rlog.encode_to(&mut self.buf);
            let checksum = crc32c::crc32c(&self.buf[origin_len..]);
            self.buf.put_u32_le(checksum);
        }
        let mut file = fs::File::create(filename)?;
        if let Some(rate_limiter) = &self.rate_limiter {
            let synced_size = file_system::write_all_with_rate_limiter(
                &mut file,
                &self.buf,
                rate_limiter,
                file_system::IoType::Compaction,
                WRITE_RLOG_BATCH_SIZE,
            )?;
            let remained_size = self.buf.len() - synced_size;
            if remained_size > 0 {
                // `write_all_with_rate_limiter` syncs only full batches. Sync
                // the remaining tail so returned metadata always points to a
                // durable rlog file.
                file.sync_data()?;
            }
        } else {
            file.write_all(&self.buf)?;
            file.sync_data()?;
        }
        let mut file = rfenginepb::RaftLogFile::default();
        file.first_index = first;
        file.last_index = last;
        if self.peer_rlog_files.is_some() {
            file.last_term = last_term;
            file.epoch_id = epoch_id;
        }

        let is_cached = self
            .rlog_cache
            .write()
            .unwrap()
            .add_rlog(peer_id, &file, &self.buf);
        Ok((file, is_cached))
    }

    fn handle_prepare_snapshot_backup(&mut self, request_token: SnapshotRequestToken) {
        let engine_id = self.manifest.get_engine_id();
        let timer = Instant::now_coarse();
        let epoch_id = self.manifest.epoch_id;
        info!("{}: start snapshot task, epoch {}", engine_id, epoch_id);
        let mut backup_meta = StoreBackupMeta::default();
        backup_meta.set_store_id(engine_id);
        backup_meta.set_epoch(epoch_id);

        let manifest = self.manifest.to_change_set(true); // Exclude tombstone peers.
        let rlog_obj_res = self.backup_raft_log_files(&manifest, &mut backup_meta);
        if let Err(rlog_obj_err) = rlog_obj_res {
            warn!(
                "{}: create snapshot rlog file failed {:?}",
                engine_id, rlog_obj_err
            );
            let _ = self.snap_tx.send(ObjectStorageTask::Snapshot {
                request_token,
                result: Err(rlog_obj_err),
            });
            return;
        }
        let rlog_obj = rlog_obj_res.unwrap();
        backup_meta.set_manifest(manifest);
        let duration = timer.saturating_elapsed();
        if rlog_obj.is_multipart() {
            info!(
                "{}: snapshot backup epoch {} prepared multipart rlog object size: {}, raft meta offset {}, takes {:?}",
                engine_id,
                epoch_id,
                rlog_obj.len(),
                backup_meta.raft_meta_start_off,
                duration,
            );
        } else {
            info!(
                "{}: snapshot backup epoch {} prepare rlog object size: {}, raft meta offset {}, takes {:?}",
                engine_id,
                epoch_id,
                rlog_obj.len(),
                backup_meta.raft_meta_start_off,
                duration,
            );
        }
        ENGINE_TAKE_SNAPSHOT_DURATION_HISTOGRAM.observe(duration.as_secs_f64());
        let prepared_snap = PreparedSnapshot {
            store_meta: backup_meta,
            rlog_obj,
        };
        let _ = self.snap_tx.send(ObjectStorageTask::Snapshot {
            request_token,
            result: Ok(prepared_snap),
        });
    }

    fn get_keyspace_id_from_peer(store_id: u64, peer_meta: &rfenginepb::PeerMeta) -> u32 {
        match utils::get_keyspace_id_from_peer(peer_meta) {
            Some(keyspace_id) => keyspace_id,
            None => {
                if peer_meta.peer_id != 0 || peer_meta.region_id != 0 {
                    warn!(
                        "{}: get unknown keyspace id peer: {:?}",
                        store_id, peer_meta
                    );
                    debug_assert!(
                        false,
                        "{}: unknown keyspace id peer: {:?}",
                        store_id, peer_meta
                    );
                }
                0
            }
        }
    }

    // Aggregate small raft-log snapshots into one local object. Large snapshots
    // upload multipart parts here and leave the final complete operation to the
    // DFS worker, so rlog is still made visible after snapshot meta.
    fn backup_raft_log_files(
        &mut self,
        manifest: &rfenginepb::ChangeSet,
        store_meta: &mut StoreBackupMeta,
    ) -> Result<PreparedSnapshotRlogObject> {
        let store_id = self.manifest.get_engine_id();
        // keyspace_id -> Vec<(peer_id, RaftLogFile)>
        let mut keyspace_map: HashMap<u32, Vec<(u64, &RaftLogFile)>> = HashMap::default();
        let mut raft_log_size: usize = 0;
        for peer in manifest.get_peers() {
            let peer_id = peer.get_peer_id();
            let keyspace_id = Self::get_keyspace_id_from_peer(store_id, peer);
            let peer_files = peer.get_files();
            let mut files = Vec::with_capacity(peer_files.len());
            for f in peer_files {
                files.push((peer_id, f));

                let rlog_size = if let Some(rlog_size) =
                    self.rlog_cache.read().unwrap().get_rlog_size(peer_id, f)
                {
                    rlog_size
                } else {
                    let file_name =
                        raft_log_file_name(&self.dir, peer_id, f.first_index, f.last_index);
                    fs::metadata(&file_name)
                        .map_err(|err| {
                            warn!("get metadata of snapshot rlog file failed";
                                "err" => ?err,
                                "file" => %file_name.display(),
                                "peer_id" => peer.get_peer_id(),
                                "region_id" => peer.get_region_id(),
                                "truncated_index" => peer.get_truncated_index(),
                                "files" => ?peer.get_files(),
                                "keyspace_id" => peer.get_keyspace_id(),
                            );
                            Error::Io(IoError::new(
                                err,
                                format!("get metadata of rlog file {}", file_name.display()),
                            ))
                        })?
                        .len() as usize
                };
                raft_log_size = raft_log_size.saturating_add(rlog_size);
            }
            keyspace_map
                .entry(keyspace_id)
                .and_modify(|k| k.append(&mut files))
                .or_insert_with(|| files);
        }

        let delayed_to_epoch_id = get_delayed_to_epoch_id(manifest);
        let object_key = snapshot_rlog_key(store_id, delayed_to_epoch_id);
        let mut object = SnapshotRlogObjectBuilder::new(
            self.snapshot_rlog_multipart_dfs.as_ref(),
            object_key,
            raft_log_size,
            self.snapshot_rlog_multipart_upload_threshold,
            self.snapshot_rlog_multipart_upload_part_size,
            self.snapshot_rlog_multipart_upload_concurrency,
            self.snapshot_rlog_max_object_size,
        )?;
        let mut rlog_meta = StoreRaftLogBackupMeta::default();
        rlog_meta.mut_header().version = 2;
        rlog_meta
            .mut_header()
            .set_compression_type(self.rlog_compression_type.to());

        let (mut total_file_cnt, mut cached_file_cnt, mut cached_file_size) = (0, 0, 0);
        for (keyspace_id, files) in keyspace_map {
            let mut keyspace_meta = KeySpaceBackupMeta::default();
            keyspace_meta.keyspace_id = keyspace_id;
            for (peer_id, file) in files {
                let rlog_cache = self.rlog_cache.read().unwrap();
                let data = if let Some(data) = rlog_cache.get_rlog_data(peer_id, file) {
                    cached_file_cnt += 1;
                    cached_file_size += data.len();
                    Cow::from(data)
                } else {
                    let file_name =
                        raft_log_file_name(&self.dir, peer_id, file.first_index, file.last_index);
                    let data = fs::read(file_name)?;
                    {
                        drop(rlog_cache);
                        self.rlog_cache
                            .write()
                            .unwrap()
                            .add_rlog(peer_id, file, &data);
                    }
                    Cow::from(data)
                };
                let compressed_data = match self.rlog_compression_type {
                    CompressionType::Lz4Compression => {
                        let mut chunk = Vec::with_capacity(data.len());
                        compress_lz4(&data, &mut chunk)?;
                        Cow::Owned(chunk)
                    }
                    CompressionType::NoCompression => Cow::Borrowed(data.as_ref()),
                };
                let mut backup_file = RaftLogBackupFile::default();
                backup_file.peer_id = peer_id;
                backup_file.start_off = object.len() as u64;
                object.put_slice(compressed_data.as_ref())?;
                backup_file.end_off = object.len() as u64;
                backup_file.first_index = file.first_index;
                backup_file.last_index = file.last_index;
                keyspace_meta.mut_files().push(backup_file);

                total_file_cnt += 1;
            }
            rlog_meta.mut_raft_logs().insert(keyspace_id, keyspace_meta);
        }

        info!("{}: backup raft log files", store_id;
            "total_file" => total_file_cnt, "total_size" => object.len(),
            "cached_file" => cached_file_cnt, "cached_size" => cached_file_size);

        store_meta.raft_meta_start_off = object.len() as u64;
        let meta = rlog_meta.write_to_bytes().unwrap();
        object.put_slice(&meta)?;

        object.finish()
    }
}

/// Returns raft logs within a target index range from a given rlog file.
///
/// - `dir`: rlog file directory
/// - `peer_id`: rlog peer id
/// - `file_first_idx`: rlog file first index
/// - `file_last_idx`: rlog file last index
/// - `start_idx`: start index of the target range (inclusive)
/// - `end_idx`: end index of the target range (inclusive)
pub(crate) fn fetch_entries_from_rlog(
    dir: &Path,
    peer_id: u64,
    file_first_idx: u64,
    file_last_idx: u64,
    start_idx: u64, // inclusive
    end_idx: u64,   // inclusive
) -> Result<Vec<RaftLogOp>> {
    if start_idx > end_idx {
        return Ok(Vec::new());
    }
    if start_idx < file_first_idx || end_idx > file_last_idx {
        return Err(Error::Other(format!(
            "peer {} index range [{}, {}] out of file range [{}, {}]",
            peer_id, start_idx, end_idx, file_first_idx, file_last_idx
        )));
    }

    // See format in Worker::write_raft_log_file
    let file = File::open(raft_log_file_name(
        dir,
        peer_id,
        file_first_idx,
        file_last_idx,
    ))?;
    let mut reader = BufReader::new(file);

    // Step 1: read rlog header.
    let mut header_buf = [0u8; RlogHeader::len()];
    reader.read_exact(&mut header_buf)?;
    let header = RlogHeader::decode(&header_buf)?;

    // Step 2: read end offsets.
    let mut end_offs_buf = vec![0u8; 4 * header.count as usize];
    reader.read_exact(&mut end_offs_buf)?;
    let end_offs: Vec<u32> = end_offs_buf
        .chunks_exact(4)
        .map(LittleEndian::read_u32)
        .collect();
    for i in 1..header.count as usize {
        debug_assert!(
            end_offs[i - 1] < end_offs[i],
            "corrupted rlog file, peer {}, file range [{},{}], \
            end offset[{}]={} > {}",
            peer_id,
            file_first_idx,
            file_last_idx,
            i - 1,
            end_offs[i - 1],
            end_offs[i]
        );
    }

    // Helper function to get the start and end offsets of a log entry.
    let entry_offsets = |log_idx: u64| -> (usize, usize) {
        let i = (log_idx - file_first_idx) as usize;
        (
            if i == 0 { 0 } else { end_offs[i - 1] as usize },
            end_offs[i] as usize,
        )
    };

    // Step 3: read and decode the raft log entries.
    // RaftLogFile format:
    // RlogHeader + [end_offsets] + [RaftLogOp] + [RaftLogOp] + ...
    let (buf_start, _) = entry_offsets(start_idx);
    reader.seek(SeekFrom::Start(
        (RlogHeader::len() + 4 * header.count as usize + buf_start) as u64,
    ))?;

    let mut result = Vec::with_capacity((end_idx - start_idx + 1) as usize);
    let mut log_buf: Vec<u8> = Vec::new();
    for log_idx in start_idx..=end_idx {
        let (log_start, log_end) = entry_offsets(log_idx);
        let len = log_end - log_start;
        log_buf.resize(len, 0);
        reader.read_exact(&mut log_buf)?;
        let (log_bytes, checksum_bytes) = log_buf.split_at(len - 4);
        let checksum = LittleEndian::read_u32(checksum_bytes);
        let actual_checksum = crc32c::crc32c(log_bytes);
        if checksum != actual_checksum {
            return Err(Error::Other(format!(
                "rlog checksum mismatch, peer {}, file range [{},{}], \
                idx {}, expected {:x}, actual {:x}",
                peer_id, file_first_idx, file_last_idx, log_idx, checksum, actual_checksum
            )));
        }
        let entry = RaftLogOp::decode(log_bytes);
        result.push(entry);
    }
    Ok(result)
}

pub(crate) fn wal_file_name(dir: &Path, epoch_id: u32, epoch_rotate_len: usize) -> PathBuf {
    let idx = epoch_id as usize % epoch_rotate_len;
    dir.join(format!("{}.wal", idx))
}

pub(crate) enum CompactTask {
    Rotate {
        epoch_id: u32,
        cache_wb_for_compact: bool,
    },
    // async write batch used for compaction.
    Write {
        wb: Arc<Vec<PeerBatch>>,
    },
    PrepareSnapshot(SnapshotRequestToken),
    Close,
}

#[derive(Default, Debug, Serialize, Deserialize)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
// If incremental is true, backup the same `epoch` WAL from the `start_offset`
pub struct BackupConfig {
    pub cluster_id: u64,
    pub store_id: u64,
    pub wal_epoch: u32,
    pub start_offset: u64,
    pub lightweight: bool,
    pub backup_ts: Option<u64>,
    pub backup_ts_wait_secs: Option<u64>,
    pub backup_ts_ttl_secs: Option<u64>,
}

// Ref `BACKUP_TS_WAIT_TIMEOUT_DEFAULT` & `BACKUP_TS_TTL_DEFAULT` in
// `native_br`.
const BACKUP_TS_WAIT_SECS_DEFAULT: u64 = 30;
const BACKUP_TS_TTL_SECS_DEFAULT: u64 = 60;

impl BackupConfig {
    pub fn backup_ts_wait_timeout(&self) -> Duration {
        Duration::from_secs(
            self.backup_ts_wait_secs
                .unwrap_or(BACKUP_TS_WAIT_SECS_DEFAULT),
        )
    }

    pub fn backup_ts_ttl(&self) -> Duration {
        Duration::from_secs(
            self.backup_ts_ttl_secs
                .unwrap_or(BACKUP_TS_TTL_SECS_DEFAULT),
        )
    }
}

pub struct BackupTask {
    pub callback: Box<dyn FnOnce(Result<StoreBackupMeta>) + Send>,
    pub(crate) file_off: u64,
    pub(crate) config: BackupConfig,
}

impl fmt::Debug for BackupTask {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("BackupTask")
            .field("file_off", &self.file_off)
            .field("config", &self.config)
            .finish()
    }
}

impl BackupTask {
    pub fn new(
        callback: Box<dyn FnOnce(Result<StoreBackupMeta>) + Send>,
        config: BackupConfig,
    ) -> Self {
        Self {
            file_off: 0,
            callback,
            config,
        }
    }

    pub fn is_lightweight(&self) -> bool {
        self.config.lightweight
    }
}

pub(crate) fn backup_callback(
    task: BackupTask,
    ret: Result<StoreBackupMeta>,
    label: &str,
    ob_start_time: Instant,
) {
    RFENGINE_BACKUP_COUNTER.with_label_values(&[label]).inc();
    if ret.is_ok() {
        RFENGINE_BACKUP_DURATION_HISTOGRAM
            .with_label_values(&[label])
            .observe(ob_start_time.saturating_elapsed_secs());
    }
    (task.callback)(ret);
}

#[cfg(test)]
mod tests {
    use std::{
        collections::HashMap,
        fs, iter,
        sync::atomic::{AtomicBool, AtomicU32, AtomicU64},
    };

    use async_trait::async_trait;
    use bytes::{Buf, Bytes};
    use kvengine::dfs::{MultipartUploadPart, Options};
    use kvproto::raft_serverpb::RegionLocalState;
    use protobuf::Message;
    use rand::{Rng, distributions::Alphanumeric};
    use rfenginepb::{ChangeSet, PeerState, StoreBackupMeta, StoreRaftLogBackupMeta};
    use tikv_util::{config::ReadableSize, defer};

    use super::*;
    use crate::{
        CompactWorker, RfEngine, RfEngineConfig, WalWriter, WriterType,
        log_batch::{RaftLogOp, RaftLogs},
        manifest::{Manifest, persist_change_set},
        raft_log_file_name, region_state_key,
        test_util::{get_txn_endkey_prefix, get_txn_startkey_prefix, init_logger},
        write_batch::PeerBatch,
    };

    const RANDOM_STR_MAX_LEN: usize = 1024;

    struct MultipartSnapshotTestDfs {
        runtime: tokio::runtime::Runtime,
    }

    impl MultipartSnapshotTestDfs {
        fn new() -> Self {
            Self {
                runtime: tokio::runtime::Builder::new_current_thread()
                    .enable_all()
                    .build()
                    .unwrap(),
            }
        }
    }

    fn unsupported_dfs_op<T>() -> kvengine::dfs::Result<T> {
        Err(kvengine::dfs::Error::Other("unsupported".to_string()))
    }

    #[async_trait]
    impl Dfs for MultipartSnapshotTestDfs {
        async fn exists(&self, _file_id: u64, _opts: Options) -> kvengine::dfs::Result<bool> {
            unsupported_dfs_op()
        }

        async fn size(&self, _file_id: u64, _opts: Options) -> kvengine::dfs::Result<u64> {
            unsupported_dfs_op()
        }

        async fn read_file(&self, _file_id: u64, _opts: Options) -> kvengine::dfs::Result<Bytes> {
            unsupported_dfs_op()
        }

        async fn create(
            &self,
            _file_id: u64,
            _data: Bytes,
            _opts: Options,
        ) -> kvengine::dfs::Result<()> {
            unsupported_dfs_op()
        }

        async fn remove(&self, _file_id: u64, _file_len: Option<u64>, _opts: Options) {}

        async fn permanently_remove(
            &self,
            _file_id: u64,
            _opts: Options,
        ) -> kvengine::dfs::Result<()> {
            unsupported_dfs_op()
        }

        fn get_runtime(&self) -> &tokio::runtime::Runtime {
            &self.runtime
        }

        fn get_prefix(&self) -> String {
            String::new()
        }

        async fn begin_upload(&self, _key: &str) -> std::result::Result<String, String> {
            Ok("upload-id".to_string())
        }

        async fn upload_part(
            &self,
            _key: &str,
            _upload_id: &str,
            part_number: u64,
            _data: Bytes,
        ) -> std::result::Result<MultipartUploadPart, String> {
            Ok(MultipartUploadPart {
                part_number,
                e_tag: format!("etag-{part_number}"),
                checksum_crc32: None,
            })
        }

        async fn complete_multipart_upload(
            &self,
            _key: &str,
            _upload_id: &str,
            _parts: Vec<MultipartUploadPart>,
        ) -> std::result::Result<(), String> {
            Ok(())
        }

        async fn abort_multipart_upload(
            &self,
            _key: &str,
            _upload_id: &str,
        ) -> std::result::Result<(), String> {
            Ok(())
        }
    }

    fn generate_random_str() -> Vec<u8> {
        let mut rng = rand::thread_rng();
        let len = rng.gen::<usize>() % RANDOM_STR_MAX_LEN + 1;
        iter::repeat(())
            .map(|()| rng.sample(Alphanumeric))
            .take(len)
            .collect()
    }

    fn new_peer_batch_with_payload(
        peer_id: u64,
        region_id: u64,
        log_count: usize,
        payload_len: usize,
    ) -> PeerBatch {
        let mut peer_batch = PeerBatch::new(peer_id, region_id, 0);
        for index in 1..=log_count {
            peer_batch.append_raft_log(RaftLogOp {
                index: index as u64,
                term: index as u32,
                e_type: 1,
                context: 0,
                data: vec![b'x'; payload_len].into(),
            });
        }
        peer_batch
    }

    fn write_keyspace_state(peer_meta: &mut rfenginepb::PeerMeta, keyspace: u32) {
        let mut state = PeerState::default();
        // mock a fake keyspace epoch version.
        let region_epoch = 10;
        state.key = region_state_key(region_epoch - 1).to_vec();
        let mut local_stat = RegionLocalState::default();
        local_stat.mut_region().start_key = get_txn_startkey_prefix(keyspace - 1).to_vec();
        local_stat.mut_region().end_key = get_txn_endkey_prefix(keyspace - 1).to_vec();
        state.value = local_stat.write_to_bytes().unwrap();
        peer_meta.mut_states().push(state.clone());

        state.key = region_state_key(region_epoch).to_vec();
        local_stat.mut_region().start_key = get_txn_startkey_prefix(keyspace).to_vec();
        local_stat.mut_region().end_key = get_txn_endkey_prefix(keyspace).to_vec();
        state.value = local_stat.write_to_bytes().unwrap();
        peer_meta.mut_states().push(state);
    }

    #[test]
    fn test_rlog_cache() {
        let peer_id = 7;

        let mut file1 = RaftLogFile::default();
        file1.first_index = 1;
        file1.last_index = 10;

        let mut file2 = RaftLogFile::default();
        file2.first_index = 1;
        file2.last_index = 20;

        let mut disabled_cache = RlogCache::new(0, 0);
        assert!(!disabled_cache.add_rlog(peer_id, &file1, b"a"));
        assert_eq!(disabled_cache.cache_size(), 0);
        assert_eq!(disabled_cache.get_rlog_data(peer_id, &file1), None);

        let mut cache = RlogCache::new(1024, 64);
        assert!(cache.add_rlog(peer_id, &file1, b"a"));
        assert!(cache.add_rlog(peer_id, &file2, b"b"));

        assert_eq!(cache.get_rlog_data(peer_id, &file1), Some(&b"a"[..]));
        assert_eq!(cache.get_rlog_data(peer_id, &file2), Some(&b"b"[..]));
    }

    #[test]
    fn test_snapshot_rlog_object_builder_reports_oversize() {
        let mut builder =
            SnapshotRlogObjectBuilder::new(None, "rlog".to_string(), 0, 0, 1, 1, 3).unwrap();

        builder.put_slice(b"abc").unwrap();
        let err = builder.put_slice(b"d").unwrap_err();

        assert!(matches!(err, Error::SnapshotOversize(4)));
    }

    #[rstest::rstest]
    #[case(false, false)]
    #[case::with_cache(true, false)]
    #[case::with_cache_and_compress(true, true)]
    fn test_backup_raft_log_files(#[case] with_cache: bool, #[case] with_compress: bool) {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let tmp_path = tmp_dir.path();
        defer!(fs::remove_dir_all(tmp_path).unwrap());
        let (_, rx) = tikv_util::mpsc::unbounded();
        let (dfs_tx, _dfs_rx) = tikv_util::mpsc::unbounded();
        let engine_id = 999;
        let manifest = Manifest::open(tmp_path, AtomicU64::new(engine_id).into()).unwrap();
        let cfg = RfEngineConfig::default();
        let mut worker = CompactWorker::new(
            tmp_path.to_path_buf(),
            &cfg,
            rx,
            dfs_tx,
            manifest,
            AtomicU32::new(0).into(),
            Arc::new(RwLock::from(RlogCache::new(0, 0))),
            CompressionType::NoCompression,
            0,
            DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE,
            1,
            MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize,
            None,
            None,
            Arc::new(AtomicBool::new(false)),
            Arc::new(ArcSwap::default()),
        );
        worker.rlog_cache = Arc::new(RwLock::from(if with_cache {
            RlogCache::new(RANDOM_STR_MAX_LEN * 80, RANDOM_STR_MAX_LEN / 2)
        } else {
            RlogCache::new(0, 0)
        }));
        worker.rlog_compression_type = if with_compress {
            CompressionType::Lz4Compression
        } else {
            CompressionType::NoCompression
        };

        let epoch = 990;
        let mut cs = ChangeSet::new();
        cs.epoch_id = epoch;
        let mut peer_rlog_files = HashMap::new();
        let mut total_file_cnt = 0;
        let mut cached_file_cnt = 0;
        for i in 100..203 {
            // peers
            let peer_id = i;
            let mut meta_pb = rfenginepb::PeerMeta::new();
            meta_pb.set_peer_id(peer_id);
            meta_pb.set_region_id(peer_id * 2);
            meta_pb.set_truncated_index(i * 10);
            let keyspace = (peer_id / 5) as u32;
            write_keyspace_state(&mut meta_pb, keyspace);
            let mut files = vec![];
            for j in 0..10 {
                // files
                let mut raft_log_file = rfenginepb::RaftLogFile::default();
                let first_index = j * 100 + 1;
                let last_index = (j + 1) * 100;
                raft_log_file.set_first_index(first_index);
                raft_log_file.set_last_index(last_index);
                let file_name = raft_log_file_name(tmp_path, peer_id, first_index, last_index);
                let content = generate_random_str();
                fs::write(file_name, &content).unwrap();
                let is_cached =
                    worker
                        .rlog_cache
                        .write()
                        .unwrap()
                        .add_rlog(peer_id, &raft_log_file, &content);
                cached_file_cnt += is_cached as usize;
                meta_pb.mut_files().push(raft_log_file);
                files.push(content);
                total_file_cnt += 1;
            }
            peer_rlog_files.insert(peer_id, files);
            cs.mut_peers().push(meta_pb);
        }
        assert!(!with_cache || cached_file_cnt > 0);
        let mut store_meta = StoreBackupMeta::default();
        let rlog_obj = worker.backup_raft_log_files(&cs, &mut store_meta).unwrap();
        let key = rlog_obj.key();
        let object = rlog_obj.as_inline_bytes().unwrap();
        let delayed_to_epoch = get_delayed_to_epoch_id(&cs);
        let raft_file_key = snapshot_rlog_key(engine_id, delayed_to_epoch);
        assert_eq!(raft_file_key, key);
        let mut raft_meta = StoreRaftLogBackupMeta::default();
        let size = object.len();
        let meta_bytes = &object.chunk()[store_meta.raft_meta_start_off as usize..size];
        raft_meta.merge_from_bytes(meta_bytes).unwrap();
        assert_eq!(raft_meta.get_header().version, 2);
        let header = raft_meta.get_header();
        if with_compress {
            assert_eq!(header.compression_type, 1);
        } else {
            assert_eq!(header.compression_type, 0);
        };

        let mut ret_file_cnt = 0;
        for (keyspace_id, keyspace_data) in raft_meta.take_raft_logs() {
            assert_eq!(keyspace_id, keyspace_data.get_keyspace_id());
            for file in keyspace_data.get_files() {
                let peer_id = file.peer_id;
                assert_eq!(keyspace_id as u64, peer_id / 5);
                let data = if with_compress {
                    decompress_lz4(&object.chunk()[file.start_off as usize..file.end_off as usize])
                        .unwrap()
                } else {
                    object.chunk()[file.start_off as usize..file.end_off as usize].to_vec()
                };
                let rlog_files = peer_rlog_files[&peer_id].clone();
                let idx = ((file.first_index - 1) / 100) as usize;
                let file_data = rlog_files[idx].clone();
                assert_eq!(
                    file_data.len(),
                    data.len(),
                    "peer id {}, index {}-{}",
                    peer_id,
                    file.first_index,
                    file.last_index
                );
                assert_eq!(
                    &file_data, &data,
                    "peer id {}, index {}-{}",
                    peer_id, file.first_index, file.last_index
                );
                ret_file_cnt += 1;
            }
        }
        assert_eq!(ret_file_cnt, total_file_cnt);
    }

    #[test]
    fn test_backup_raft_log_files_reports_missing_rlog_path() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let tmp_path = tmp_dir.path();
        defer!(fs::remove_dir_all(tmp_path).unwrap());
        let (_, rx) = tikv_util::mpsc::unbounded();
        let (dfs_tx, _dfs_rx) = tikv_util::mpsc::unbounded();
        let engine_id = 999;
        let manifest = Manifest::open(tmp_path, AtomicU64::new(engine_id).into()).unwrap();
        let cfg = RfEngineConfig::default();
        let mut worker = CompactWorker::new(
            tmp_path.to_path_buf(),
            &cfg,
            rx,
            dfs_tx,
            manifest,
            AtomicU32::new(0).into(),
            Arc::new(RwLock::from(RlogCache::new(0, 0))),
            CompressionType::NoCompression,
            0,
            DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE,
            1,
            MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize,
            None,
            None,
            Arc::new(AtomicBool::new(false)),
            Arc::new(ArcSwap::default()),
        );

        let peer_id = 7;
        let first_index = 1;
        let last_index = 10;
        let missing_file = raft_log_file_name(tmp_path, peer_id, first_index, last_index);
        let mut raft_log_file = rfenginepb::RaftLogFile::default();
        raft_log_file.set_first_index(first_index);
        raft_log_file.set_last_index(last_index);

        let mut peer = rfenginepb::PeerMeta::new();
        peer.set_peer_id(peer_id);
        peer.set_region_id(peer_id * 2);
        peer.set_truncated_index(first_index);
        peer.set_keyspace_id(1);
        write_keyspace_state(&mut peer, 1);
        peer.mut_files().push(raft_log_file);

        let mut change_set = ChangeSet::new();
        change_set.mut_peers().push(peer);

        let mut store_meta = StoreBackupMeta::default();
        let err = match worker.backup_raft_log_files(&change_set, &mut store_meta) {
            Ok(_) => panic!("backup_raft_log_files should report missing rlog path"),
            Err(err) => err,
        };
        match err {
            Error::Io(io_err) => {
                assert_eq!(io_err.kind(), std::io::ErrorKind::NotFound);
                assert!(
                    io_err.ctx().contains("get metadata of rlog file"),
                    "unexpected io context: {}",
                    io_err.ctx()
                );
                assert!(
                    io_err.ctx().contains(&missing_file.display().to_string()),
                    "io context should include missing file {}, got {}",
                    missing_file.display(),
                    io_err.ctx()
                );
            }
            other => panic!("unexpected error: {:?}", other),
        }
    }

    #[test]
    fn test_backup_raft_log_files_returns_multipart_at_threshold() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let tmp_path = tmp_dir.path();
        defer!(fs::remove_dir_all(tmp_path).unwrap());
        let (_, rx) = tikv_util::mpsc::unbounded();
        let (dfs_tx, _dfs_rx) = tikv_util::mpsc::unbounded();
        let engine_id = 999;
        let manifest = Manifest::open(tmp_path, AtomicU64::new(engine_id).into()).unwrap();
        let snapshot_dfs: Arc<dyn Dfs> = Arc::new(MultipartSnapshotTestDfs::new());
        let cfg = RfEngineConfig::default();
        let mut worker = CompactWorker::new(
            tmp_path.to_path_buf(),
            &cfg,
            rx,
            dfs_tx,
            manifest,
            AtomicU32::new(0).into(),
            Arc::new(RwLock::from(RlogCache::new(0, 0))),
            CompressionType::NoCompression,
            1,
            2,
            1,
            MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize,
            Some(snapshot_dfs),
            None,
            Arc::new(AtomicBool::new(false)),
            Arc::new(ArcSwap::default()),
        );

        let peer_id = 7;
        let first_index = 1;
        let last_index = 10;
        let mut raft_log_file = rfenginepb::RaftLogFile::default();
        raft_log_file.set_first_index(first_index);
        raft_log_file.set_last_index(last_index);
        fs::write(
            raft_log_file_name(tmp_path, peer_id, first_index, last_index),
            b"rlog",
        )
        .unwrap();

        let mut peer = rfenginepb::PeerMeta::new();
        peer.set_peer_id(peer_id);
        peer.set_region_id(peer_id * 2);
        peer.set_truncated_index(first_index);
        peer.set_keyspace_id(1);
        write_keyspace_state(&mut peer, 1);
        peer.mut_files().push(raft_log_file);

        let mut change_set = ChangeSet::new();
        change_set.mut_peers().push(peer);

        let mut store_meta = StoreBackupMeta::default();
        let rlog_obj = worker
            .backup_raft_log_files(&change_set, &mut store_meta)
            .unwrap();

        assert!(rlog_obj.is_multipart());
        assert!(rlog_obj.as_inline_bytes().is_none());
        assert!(store_meta.raft_meta_start_off > 0);
        assert!(rlog_obj.len() > store_meta.raft_meta_start_off as usize);
    }

    fn generate_rlog_files(
        worker: &mut CompactWorker,
        raft_logs: &mut RaftLogs,
        peer_id: u64,
        region_id: u64,
        keyspace_id: u32,
        first_index: u64,
        last_index: u64,
    ) -> (rfenginepb::RaftLogFile, bool /* is_cached */) {
        let mut peer_batch = PeerBatch::new(peer_id, region_id, keyspace_id);
        for index in first_index..=last_index {
            let op = RaftLogOp {
                index,
                term: index as u32,
                e_type: 1,
                context: 2,
                data: generate_random_str().into(),
            };
            peer_batch.append_raft_log(op.clone());
            raft_logs.append(peer_id, op);
        }
        let rlog_epoch = 3;
        let files = worker.write_raft_log_files(peer_batch, rlog_epoch).unwrap();
        assert_eq!(files.len(), 1);
        let (file, is_cached) = files.into_iter().next().unwrap();
        if worker.peer_rlog_files.is_some() {
            assert_eq!(file.last_term, last_index as u32);
            assert_eq!(file.epoch_id, rlog_epoch);
        }
        (file, is_cached)
    }

    #[rstest::rstest]
    #[case(false, false)]
    #[case::with_cache(true, false)]
    #[case::with_cache_and_compress(true, true)]
    fn test_backup_and_load_raft_log_files(#[case] with_cache: bool, #[case] with_compress: bool) {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let tmp_path = tmp_dir.path();
        defer!(fs::remove_dir_all(tmp_path).unwrap());
        let (_, rx) = tikv_util::mpsc::unbounded();
        let (dfs_tx, _dfs_rx) = tikv_util::mpsc::unbounded();
        let engine_id = 1999;
        let manifest = Manifest::open(tmp_path, AtomicU64::new(engine_id).into()).unwrap();
        let cfg = RfEngineConfig::default();
        let mut worker = CompactWorker::new(
            tmp_path.to_path_buf(),
            &cfg,
            rx,
            dfs_tx,
            manifest,
            AtomicU32::new(0).into(),
            Arc::new(RwLock::from(RlogCache::new(0, 0))),
            CompressionType::NoCompression,
            0,
            DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE,
            1,
            MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize,
            None,
            None,
            Arc::new(AtomicBool::new(false)),
            Arc::new(ArcSwap::default()),
        );
        worker.rlog_cache = Arc::new(RwLock::from(if with_cache {
            RlogCache::new(RANDOM_STR_MAX_LEN * 100 * 5, RANDOM_STR_MAX_LEN * 100 / 2)
        } else {
            RlogCache::new(0, 0)
        }));
        worker.rlog_compression_type = if with_compress {
            CompressionType::Lz4Compression
        } else {
            CompressionType::NoCompression
        };

        let mut cs = ChangeSet::new();
        let mut peer_data_map = HashMap::new();
        let peers_range = 100..150;
        let mut total_file_cnt = 0;
        let mut cached_file_cnt = 0;
        for i in peers_range.clone() {
            let peer_id = i;
            let region_id = peer_id * 2;
            let mut meta_pb = rfenginepb::PeerMeta::new();
            meta_pb.set_peer_id(peer_id);
            meta_pb.set_region_id(region_id);
            meta_pb.set_truncated_index(0);
            let keyspace_id = i as u32 / 10;
            write_keyspace_state(&mut meta_pb, keyspace_id);
            let mut peer_raft_log = RaftLogs::default();
            for j in 0..10 {
                let first_index = j * 100 + 1;
                let last_index = (j + 1) * 100;
                let (raft_log_file, is_cached) = generate_rlog_files(
                    &mut worker,
                    &mut peer_raft_log,
                    peer_id,
                    region_id,
                    1,
                    first_index,
                    last_index,
                );
                meta_pb.mut_files().push(raft_log_file);
                total_file_cnt += 1;
                cached_file_cnt += is_cached as usize;
            }
            cs.mut_peers().push(meta_pb);
            cs.epoch_id += 1;
            peer_data_map.insert(peer_id, peer_raft_log);
        }
        assert!(!with_cache || cached_file_cnt > 0);

        // 1. backup
        let mut store_meta = StoreBackupMeta::default();
        let rlog_obj = worker.backup_raft_log_files(&cs, &mut store_meta).unwrap();
        let key = rlog_obj.key();
        let object = rlog_obj.as_inline_bytes().unwrap();
        let delayed_to_epoch_id = get_delayed_to_epoch_id(&cs);
        let raft_file_key = snapshot_rlog_key(engine_id, delayed_to_epoch_id);
        assert_eq!(raft_file_key, key);

        let mut raft_meta = StoreRaftLogBackupMeta::default();
        let size = object.len();
        let meta_bytes = &object.chunk()[store_meta.raft_meta_start_off as usize..size];
        raft_meta.merge_from_bytes(meta_bytes).unwrap();
        assert_eq!(raft_meta.get_header().version, 2);
        let header = raft_meta.get_header();
        if with_compress {
            assert_eq!(header.compression_type, 1);
        } else {
            assert_eq!(header.compression_type, 0);
        };
        // 2. restore to new dir
        let tmp_dir2 = tempfile::tempdir().unwrap();
        let tmp_path2 = tmp_dir2.path();
        defer!(fs::remove_dir_all(tmp_path2).unwrap());
        // restore raft log files by keyspace
        let mut restore_file_cnt = 0;
        for (keyspace_id, keyspace_data) in raft_meta.take_raft_logs() {
            assert_eq!(keyspace_id, keyspace_data.get_keyspace_id());
            for file in keyspace_data.get_files() {
                let peer_id = file.peer_id;
                assert_eq!(peer_id / 10, keyspace_id as u64);
                let data = if with_compress {
                    decompress_lz4(&object.chunk()[file.start_off as usize..file.end_off as usize])
                        .unwrap()
                } else {
                    object.chunk()[file.start_off as usize..file.end_off as usize].to_vec()
                };
                let file_name =
                    raft_log_file_name(tmp_path2, peer_id, file.first_index, file.last_index);
                fs::write(file_name, data).unwrap();
                restore_file_cnt += 1;
            }
        }
        assert_eq!(total_file_cnt, restore_file_cnt);
        // restore manifest file
        let manifest = Manifest::open(tmp_path2, AtomicU64::new(engine_id).into()).unwrap();
        persist_change_set(&manifest.file, 0, &cs).unwrap();

        // 3. open RfEngine with restored dir.
        let wal_size = 4 * 1024 * 1024;
        let cfg = RfEngineConfig::new(wal_size);
        // Hack wal file to let RfEngine::open pass.
        let mut wal_writer = WalWriter::new(
            tmp_path2,
            &cfg,
            AtomicU32::new(cs.epoch_id + 1).into(),
            WriterType::Sync,
        );
        wal_writer.open_file(cs.epoch_id + 1, 0).unwrap();
        // checksum inner should succeeds.
        let engine = RfEngine::open(tmp_path2, &cfg, None, None).unwrap();
        // 4. Check peer data, restored data should be same with previous one.
        for peer_id in peers_range {
            let cache_peer_data = peer_data_map.get(&peer_id).unwrap();
            for index in 1..=1000 {
                let entry1 = cache_peer_data.get(index);
                let entry2 = engine.get_raft_entry(peer_id, index);
                assert_eq!(entry1, entry2);
            }
        }
    }

    #[test]
    fn test_get_keyspace_id_from_peer() {
        init_logger();
        let mut peer_meta = rfenginepb::PeerMeta::default();
        assert_eq!(CompactWorker::get_keyspace_id_from_peer(1, &peer_meta), 0);
        write_keyspace_state(&mut peer_meta, 100);
        write_keyspace_state(&mut peer_meta, 200);
        assert_eq!(CompactWorker::get_keyspace_id_from_peer(1, &peer_meta), 200);
    }

    #[test]
    fn test_build_sub_batches_balances_skewed_peer_sizes() {
        let heavy_peers =
            (1..=3).map(|peer_id| new_peer_batch_with_payload(peer_id, peer_id * 10, 64, 512));
        let small_peers =
            (101..=103).map(|peer_id| new_peer_batch_with_payload(peer_id, peer_id * 10, 2, 32));
        let peer_batches = heavy_peers.chain(small_peers).collect();

        let sub_batches = build_sub_batches_for_compaction(peer_batches, 2);

        assert_eq!(sub_batches.len(), 3);
        for sub_batch in &sub_batches {
            assert!(sub_batch.len() <= 2);
            assert_eq!(
                sub_batch
                    .iter()
                    .filter(|peer_batch| peer_batch.peer_id <= 3)
                    .count(),
                1
            );
        }
    }

    #[test]
    fn test_build_sub_batches_trims_logs_covered_by_truncated_idx() {
        let mut peer_batch = new_peer_batch_with_payload(1, 10, 8, 16);
        peer_batch.meta.truncated_idx = 5;

        let sub_batches = build_sub_batches_for_compaction(vec![peer_batch], 2);

        assert_eq!(sub_batches.len(), 1);
        assert_eq!(sub_batches[0].len(), 1);
        let trimmed = &sub_batches[0][0];
        assert_eq!(trimmed.truncated_idx, 5);
        assert_eq!(trimmed.raft_logs.front().map(|log| log.index), Some(6));
        assert_eq!(trimmed.raft_logs.back().map(|log| log.index), Some(8));
        assert_eq!(trimmed.raft_logs.len(), 3);
    }

    #[test]
    fn test_raft_log_file_size_limit() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let tmp_path = tmp_dir.path();
        defer!(fs::remove_dir_all(tmp_path).unwrap());
        let (_, rx) = tikv_util::mpsc::unbounded();
        let (dfs_tx, _dfs_rx) = tikv_util::mpsc::unbounded();

        let manifest = Manifest::open(tmp_path, AtomicU64::new(1).into()).unwrap();
        let mut cfg = RfEngineConfig::default();
        cfg.rlog_file_size = ReadableSize::kb(10);
        let mut worker = CompactWorker::new(
            tmp_path.to_path_buf(),
            &cfg,
            rx,
            dfs_tx,
            manifest,
            AtomicU32::new(0).into(),
            Arc::new(RwLock::from(RlogCache::new(0, 0))),
            CompressionType::NoCompression,
            0,
            DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE,
            1,
            MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize,
            None,
            None,
            Arc::new(AtomicBool::new(false)),
            Arc::new(ArcSwap::default()),
        );

        let mut peer_batch = PeerBatch::new(1, 1000, 1);
        for index in 1..=30 {
            let op = RaftLogOp {
                index,
                term: index as u32,
                e_type: 1,
                context: 2,
                data: "a".repeat(1000).into(), // 1000 bytes
            };
            peer_batch.append_raft_log(op.clone());
        }
        let rlog_epoch = 3;
        let files = worker.write_raft_log_files(peer_batch, rlog_epoch).unwrap();
        // Check that there should be 3 rlog files, each with 10 entries.
        assert_eq!(files.len(), 3);
        for (i, (file, _)) in files.iter().enumerate() {
            assert_eq!(file.first_index, (i as u64) * 10 + 1);
            assert_eq!(file.last_index, (i as u64 + 1) * 10);
            if worker.peer_rlog_files.is_some() {
                assert_eq!(file.last_term, file.last_index as u32);
                assert_eq!(file.epoch_id, rlog_epoch);
            }
        }

        // Test that if a single raft entry exceeds the size limit, an
        // error will be returned.
        let mut peer_batch = PeerBatch::new(1, 1000, 1);
        peer_batch.append_raft_log(RaftLogOp {
            index: 1,
            term: 1,
            e_type: 1,
            context: 2,
            data: "a".repeat(20 * 1024).into(), // 20KB
        });
        let files = worker.write_raft_log_files(peer_batch, 4).unwrap();
        assert_eq!(files.len(), 1);
    }

    #[test]
    fn test_fetch_entries_from_rlog() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let tmp_path = tmp_dir.path();
        defer!(fs::remove_dir_all(tmp_path).unwrap());
        let (_, rx) = tikv_util::mpsc::unbounded();
        let (dfs_tx, _dfs_rx) = tikv_util::mpsc::unbounded();

        let manifest = Manifest::open(tmp_path, AtomicU64::new(1).into()).unwrap();
        let cfg = RfEngineConfig::default();
        let mut worker = CompactWorker::new(
            tmp_path.to_path_buf(),
            &cfg,
            rx,
            dfs_tx,
            manifest,
            AtomicU32::new(0).into(),
            Arc::new(RwLock::from(RlogCache::new(0, 0))),
            CompressionType::NoCompression,
            0,
            DEFAULT_SNAPSHOT_RLOG_MULTIPART_UPLOAD_PART_SIZE,
            1,
            MAX_SNAPSHOT_RLOG_OBJECT_SIZE as usize,
            None,
            None,
            Arc::new(AtomicBool::new(false)),
            Arc::new(ArcSwap::default()),
        );
        let peer_id = 1001;
        let (raft_log_file, _) = generate_rlog_files(
            &mut worker,
            &mut RaftLogs::default(),
            peer_id,
            1000, // region id
            1,    // keyspace id
            10,   // first_index
            20,   // last_index
        );

        for (start_idx, end_idx, expected_len) in vec![
            // Success cases
            (11, 10, Ok(0)),
            (10, 10, Ok(1)),
            (15, 15, Ok(1)),
            (20, 20, Ok(1)),
            (15, 19, Ok(5)),
            (10, 15, Ok(6)),
            (10, 20, Ok(11)),
            // Failure cases
            (9, 9, Err(())),
            (21, 21, Err(())),
            (9, 10, Err(())),
            (10, 21, Err(())),
            (9, 21, Err(())),
        ] {
            let result = fetch_entries_from_rlog(
                tmp_path,
                peer_id,
                raft_log_file.first_index,
                raft_log_file.last_index,
                start_idx,
                end_idx,
            );
            match expected_len {
                Ok(expected_len) => {
                    let logs = result.expect("expected success");
                    assert_eq!(logs.len(), expected_len);
                    for (i, entry) in logs.iter().enumerate() {
                        let expected_idx = start_idx + i as u64;
                        assert_eq!(entry.index, expected_idx);
                        assert_eq!(entry.term, expected_idx as u32);
                    }
                }
                Err(_) => {
                    result.unwrap_err();
                }
            }
        }
    }
}