Repository: tikv/tikv (language: rust)

Issue:
Snapshot generation speed should be limited by actual I/O usage
## Development Task
Snapshot generation speed should be limited by actual I/O usage instead of the logical size of scanned key-value pairs. The actual I/O usage is based on the compressed SST file size, which is typically 2x smaller than the logical size.


Discussion / hints:

Fix patch:
diff --git a/components/file_system/src/io_stats/mod.rs b/components/file_system/src/io_stats/mod.rs
index 9a2bc242ed9..7a5551e5044 100644
--- a/components/file_system/src/io_stats/mod.rs
+++ b/components/file_system/src/io_stats/mod.rs
@@ -31,7 +31,7 @@ mod stub {
     }
 
     pub fn get_thread_io_bytes_total() -> Result<IoBytes, String> {
-        Err("unimplemented".into())
+        Ok(IoBytes::default())
     }
 }
 #[cfg(not(any(target_os = "linux", feature = "bcc-iosnoop")))]
diff --git a/components/raftstore/src/store/snap.rs b/components/raftstore/src/store/snap.rs
index 8516703de57..20b9fba22a1 100644
--- a/components/raftstore/src/store/snap.rs
+++ b/components/raftstore/src/store/snap.rs
@@ -899,6 +899,7 @@ impl Snapshot {
                         .get_actual_max_per_file_size(allow_multi_files_snapshot),
                     &self.mgr.limiter,
                     self.mgr.encryption_key_manager.clone(),
+                    for_balance,
                 )?
             };
             SNAPSHOT_LIMIT_GENERATE_BYTES.inc_by(cf_stat.total_size as u64);
@@ -2593,6 +2594,18 @@ pub mod tests {
         db_opt: Option<DbOptions>,
         cf_opts: Option<Vec<(&'static str, CfOptions)>>,
     ) -> Result<E>
+    where
+        E: KvEngine + KvEngineConstructorExt,
+    {
+        open_test_db_with_nkeys(path, db_opt, cf_opts, 100)
+    }
+
+    pub fn open_test_db_with_nkeys<E>(
+        path: &Path,
+        db_opt: Option<DbOptions>,
+        cf_opts: Option<Vec<(&'static str, CfOptions)>>,
+        nkeys: u64,
+    ) -> Result<E>
     where
         E: KvEngine + KvEngineConstructorExt,
     {
@@ -2602,7 +2615,7 @@ pub mod tests {
             let mut p = Peer::default();
             p.set_store_id(TEST_STORE_ID);
             p.set_id((i + 1) as u64);
-            for k in 0..100 {
+            for k in 0..nkeys {
                 let key = keys::data_key(format!("akey{}", k).as_bytes());
                 db.put_msg_cf(cf, &key[..], &p)?;
             }
diff --git a/components/raftstore/src/store/snap/io.rs b/components/raftstore/src/store/snap/io.rs
index 2962dd903e6..1d1597e95e2 100644
--- a/components/raftstore/src/store/snap/io.rs
+++ b/components/raftstore/src/store/snap/io.rs
@@ -13,7 +13,9 @@ use engine_traits::{
     SstCompressionType, SstReader, SstWriter, SstWriterBuilder, WriteBatch,
 };
 use fail::fail_point;
-use file_system::{File, OpenOptions};
+#[cfg(not(test))]
+use file_system::get_thread_io_bytes_total;
+use file_system::{File, IoBytes, IoType, OpenOptions, WithIoType};
 use kvproto::encryptionpb::EncryptionMethod;
 use tikv_util::{
     box_try,
@@ -24,6 +26,11 @@ use tikv_util::{
 
 use super::{CfFile, Error, IO_LIMITER_CHUNK_SIZE};
 
+// This defines the number of bytes scanned before trigger an I/O limiter check.
+// It is used instead of checking the I/O limiter for each scan to reduce cpu
+// overhead.
+const SCAN_BYTES_PER_IO_LIMIT_CHECK: usize = 8 * 1024;
+
 /// Used to check a procedure is stale or not.
 pub trait StaleDetector {
     fn is_stale(&self) -> bool;
@@ -104,6 +111,25 @@ where
     Ok(stats)
 }
 
+#[cfg(not(test))]
+fn get_thread_io_bytes_stats() -> Result<file_system::IoBytes, String> {
+    get_thread_io_bytes_total()
+}
+
+#[cfg(test)]
+fn get_thread_io_bytes_stats() -> Result<file_system::IoBytes, String> {
+    use std::cell::Cell;
+    thread_local! {
+        static TOTAL_BYTES: Cell<IoBytes> = Cell::new(IoBytes::default());
+    }
+    let mut new_bytes = TOTAL_BYTES.get();
+    // We use 2 as the factor because the compression ratio of SST files with the
+    // zstd algorithm is empirically around 2x.
+    new_bytes.read += SCAN_BYTES_PER_IO_LIMIT_CHECK as u64 / 2;
+    TOTAL_BYTES.set(new_bytes);
+    Ok(new_bytes)
+}
+
 /// Build a snapshot file for the given column family in sst format.
 /// If there are no key-value pairs fetched, no files will be created at `path`,
 /// otherwise the file will be created and synchronized.
@@ -116,6 +142,7 @@ pub fn build_sst_cf_file_list<E>(
     raw_size_per_file: u64,
     io_limiter: &Limiter,
     key_mgr: Option<Arc<DataKeyManager>>,
+    for_balance: bool,
 ) -> Result<BuildStatistics, Error>
 where
     E: KvEngine,
@@ -176,6 +203,25 @@ where
     };
 
     let instant = Instant::now();
+    let _io_type_guard = WithIoType::new(if for_balance {
+        IoType::LoadBalance
+    } else {
+        IoType::Replication
+    });
+    let mut prev_io_bytes = get_thread_io_bytes_stats().unwrap();
+    let mut next_io_check_size = stats.total_size + SCAN_BYTES_PER_IO_LIMIT_CHECK;
+    let handle_read_io_usage = |prev_io_bytes: &mut IoBytes, remained_quota: &mut usize| {
+        let cur_io_bytes = get_thread_io_bytes_stats().unwrap();
+        let read_delta = (cur_io_bytes.read - prev_io_bytes.read) as usize;
+
+        while read_delta > *remained_quota {
+            io_limiter.blocking_consume(IO_LIMITER_CHUNK_SIZE);
+            *remained_quota += IO_LIMITER_CHUNK_SIZE;
+        }
+        *remained_quota -= read_delta;
+        *prev_io_bytes = cur_io_bytes;
+    };
+
     box_try!(snap.scan(cf, start_key, end_key, false, |key, value| {
         let entry_len = key.len() + value.len();
         if file_length + entry_len > raw_size_per_file as usize {
@@ -202,15 +248,16 @@ where
             }
         }
 
-        while entry_len > remained_quota {
-            // It's possible to acquire more than necessary, but let it be.
-            io_limiter.blocking_consume(IO_LIMITER_CHUNK_SIZE);
-            remained_quota += IO_LIMITER_CHUNK_SIZE;
-        }
-        remained_quota -= entry_len;
-
         stats.key_count += 1;
         stats.total_size += entry_len;
+
+        if stats.total_size >= next_io_check_size {
+            // TODO(@hhwyt): Consider incorporating snapshot file write I/O into the
+            // limiting mechanism.
+            handle_read_io_usage(&mut prev_io_bytes, &mut remained_quota);
+            next_io_check_size = stats.total_size + SCAN_BYTES_PER_IO_LIMIT_CHECK;
+        }
+
         if let Err(e) = sst_writer.borrow_mut().put(key, value) {
             let io_error = io::Error::new(io::ErrorKind::Other, e);
             return Err(io_error.into());
@@ -218,6 +265,10 @@ where
         file_length += entry_len;
         Ok(true)
     }));
+    // Handle the IO generated by the remaining key-value pairs less than
+    // SCAN_BYTES_PER_IO_LIMIT_CHECK.
+    handle_read_io_usage(&mut prev_io_bytes, &mut remained_quota);
+
     if stats.key_count > 0 {
         box_try!(finish_sst_writer(sst_writer.into_inner(), path, key_mgr));
         cf_file.add_file(file_id);
@@ -566,6 +617,7 @@ mod tests {
                         *max_file_size,
                         &limiter,
                         db_opt.as_ref().and_then(|opt| opt.get_key_manager()),
+                        true,
                     )
                     .unwrap();
                     if stats.key_count == 0 {
@@ -602,4 +654,48 @@ mod tests {
             }
         }
     }
+
+    // This test verifies that building SST files is effectively limited by the I/O
+    // limiter based on actual I/O usage. It achieve this by adding an I/O limiter
+    // and asserting that the elapsed time for building SST files exceeds the
+    // lower bound enforced by the I/O limiter.
+    //
+    // In this test, the I/O limiter is configured with a throughput limit 8000
+    // bytes/sec. A dataset of 1000 keys (totaling 11, 890 bytes) is generated  to
+    // trigger two I/O limiter checks, as the default SCAN_BYTES_PER_IO_LIMIT_CHECK
+    // is 8192 bytes. During each check, the mocked `get_thread_io_bytes_stats`
+    // function returns 4096 bytes of I/O usage, resulting in total of 8192 bytes.
+    // With the 8000 bytes/sec limitation, we assert that the elapsed time must
+    // exceed 1 second.
+    #[test]
+    fn test_build_sst_with_io_limiter() {
+        let dir = Builder::new().prefix("test-io-limiter").tempdir().unwrap();
+        let db = open_test_db_with_nkeys(dir.path(), None, None, 1000).unwrap();
+        let bytes_per_sec = 8000_f64;
+        let limiter = Limiter::new(bytes_per_sec);
+        let snap_dir = Builder::new().prefix("snap-dir").tempdir().unwrap();
+        let mut cf_file = CfFile {
+            cf: CF_DEFAULT,
+            path: PathBuf::from(snap_dir.path()),
+            file_prefix: "test_sst".to_string(),
+            file_suffix: SST_FILE_SUFFIX.to_string(),
+            ..Default::default()
+        };
+
+        let start = Instant::now();
+        let stats = build_sst_cf_file_list::<KvTestEngine>(
+            &mut cf_file,
+            &db,
+            &db.snapshot(),
+            &keys::data_key(b""),
+            &keys::data_key(b"z"),
+            u64::MAX,
+            &limiter,
+            None,
+            true,
+        )
+        .unwrap();
+        assert_eq!(stats.total_size, 11890);
+        assert!(start.saturating_elapsed_secs() > 1_f64);
+    }
 }


Test patch:
diff --git a/tests/integrations/storage/test_titan.rs b/tests/integrations/storage/test_titan.rs
index 86bcb609426..a99f0383c81 100644
--- a/tests/integrations/storage/test_titan.rs
+++ b/tests/integrations/storage/test_titan.rs
@@ -380,6 +380,7 @@ fn test_delete_files_in_range_for_titan() {
         u64::MAX,
         &limiter,
         None,
+        true,
     )
     .unwrap();
     let mut cf_file_write = CfFile::new(
@@ -397,6 +398,7 @@ fn test_delete_files_in_range_for_titan() {
         u64::MAX,
         &limiter,
         None,
+        true,
     )
     .unwrap();