import: sst_importer support download SST and rewrite into keyspace data. (tikv#14046)

ref tikv#12999

import: sst_importer support download SST and rewrite into keyspace data.

Signed-off-by: iosmanthus <[email protected]>

Author: iosmanthus

Cargo.lock

@@ -10,6 +10,7 @@ kvproto = { workspace = true }
 log_wrappers = { workspace = true }
 thiserror = "1.0"
 tikv_alloc = { workspace = true }
+tikv_util = { workspace = true }
 
 [dev-dependencies]
 panic_hook = { workspace = true }

components/keys/Cargo.toml

@@ -6,11 +6,21 @@
 
 use std::ops::Bound::{self, *};
 
+use tikv_util::codec::bytes::encode_bytes;
+
 /// An error indicating the key cannot be rewritten because it does not start
 /// with the given prefix.
 #[derive(PartialEq, Debug, Clone)]
 pub struct WrongPrefix;
 
+pub fn encode_bound(bound: Bound<Vec<u8>>) -> Bound<Vec<u8>> {
+    match bound {
+        Included(k) => Included(encode_bytes(&k)),
+        Excluded(k) => Excluded(encode_bytes(&k)),
+        Unbounded => Unbounded,
+    }
+}
+
 /// Rewrites the prefix of a byte array.
 pub fn rewrite_prefix(
     old_prefix: &[u8],

components/keys/src/rewrite.rs

@@ -15,6 +15,7 @@ use engine_traits::{
     iter_option, EncryptionKeyManager, Iterator, KvEngine, RefIterable, SstMetaInfo, SstReader,
 };
 use file_system::{get_io_rate_limiter, sync_dir, File, OpenOptions};
+use keys::data_key;
 use kvproto::{import_sstpb::*, kvrpcpb::ApiVersion};
 use tikv_util::time::Instant;
 use uuid::{Builder as UuidBuilder, Uuid};
@@ -336,7 +337,7 @@ impl ImportDir {
                     let sst_reader = RocksSstReader::open_with_env(path_str, Some(env))?;
 
                     for &(start, end) in TIDB_RANGES_COMPLEMENT {
-                        let opt = iter_option(start, end, false);
+                        let opt = iter_option(&data_key(start), &data_key(end), false);
                         let mut iter = sst_reader.iter(opt)?;
                         if iter.seek(start)? {
                             error!(

components/sst_importer/src/import_file.rs

@@ -32,7 +32,10 @@ use kvproto::{
     kvrpcpb::ApiVersion,
 };
 use tikv_util::{
-    codec::stream_event::{EventEncoder, EventIterator, Iterator as EIterator},
+    codec::{
+        bytes::{decode_bytes_in_place, encode_bytes},
+        stream_event::{EventEncoder, EventIterator, Iterator as EIterator},
+    },
     config::ReadableSize,
     stream::block_on_external_io,
     sys::SysQuota,
@@ -53,13 +56,18 @@ use crate::{
 #[derive(Default, Debug, Clone)]
 pub struct DownloadExt<'a> {
     cache_key: Option<&'a str>,
+    req_type: DownloadRequestType,
 }
 
 impl<'a> DownloadExt<'a> {
-    pub fn cache_key(self, key: &'a str) -> Self {
-        Self {
-            cache_key: Some(key),
-        }
+    pub fn cache_key(mut self, key: &'a str) -> Self {
+        self.cache_key = Some(key);
+        self
+    }
+
+    pub fn req_type(mut self, req_type: DownloadRequestType) -> Self {
+        self.req_type = req_type;
+        self
     }
 }
 
@@ -896,16 +904,20 @@ impl SstImporter {
         let sst_reader = RocksSstReader::open_with_env(dst_file_name, Some(env))?;
         sst_reader.verify_checksum()?;
 
+        // undo key rewrite so we could compare with the keys inside SST
+        let old_prefix = rewrite_rule.get_old_key_prefix();
+        let new_prefix = rewrite_rule.get_new_key_prefix();
+        let req_type = ext.req_type;
+
         debug!("downloaded file and verified";
             "meta" => ?meta,
             "name" => name,
             "path" => dst_file_name,
+            "old_prefix" => log_wrappers::Value::key(old_prefix),
+            "new_prefix" => log_wrappers::Value::key(new_prefix),
+            "req_type" => ?req_type,
         );
 
-        // undo key rewrite so we could compare with the keys inside SST
-        let old_prefix = rewrite_rule.get_old_key_prefix();
-        let new_prefix = rewrite_rule.get_new_key_prefix();
-
         let range_start = meta.get_range().get_start();
         let range_end = meta.get_range().get_end();
         let range_start_bound = key_to_bound(range_start);
@@ -915,21 +927,26 @@ impl SstImporter {
             key_to_bound(range_end)
         };
 
-        let range_start =
+        let mut range_start =
             keys::rewrite::rewrite_prefix_of_start_bound(new_prefix, old_prefix, range_start_bound)
                 .map_err(|_| Error::WrongKeyPrefix {
                     what: "SST start range",
                     key: range_start.to_vec(),
                     prefix: new_prefix.to_vec(),
                 })?;
-        let range_end =
+        let mut range_end =
             keys::rewrite::rewrite_prefix_of_end_bound(new_prefix, old_prefix, range_end_bound)
                 .map_err(|_| Error::WrongKeyPrefix {
                     what: "SST end range",
                     key: range_end.to_vec(),
                     prefix: new_prefix.to_vec(),
                 })?;
 
+        if req_type == DownloadRequestType::Keyspace {
+            range_start = keys::rewrite::encode_bound(range_start);
+            range_end = keys::rewrite::encode_bound(range_end);
+        }
+
         let start_rename_rewrite = Instant::now();
         // read the first and last keys from the SST, determine if we could
         // simply move the entire SST instead of iterating and generate a new one.
@@ -942,9 +959,15 @@ impl SstImporter {
                 return Ok(None);
             }
             if !iter.seek_to_first()? {
+                let mut range = meta.get_range().clone();
+                if req_type == DownloadRequestType::Keyspace {
+                    *range.mut_start() = encode_bytes(&range.take_start());
+                    *range.mut_end() = encode_bytes(&range.take_end());
+                }
                 // the SST is empty, so no need to iterate at all (should be impossible?)
-                return Ok(Some(meta.get_range().clone()));
+                return Ok(Some(range));
             }
+
             let start_key = keys::origin_key(iter.key());
             if is_before_start_bound(start_key, &range_start) {
                 // SST's start is before the range to consume, so needs to iterate to skip over
@@ -995,8 +1018,10 @@ impl SstImporter {
         }
 
         // perform iteration and key rewrite.
-        let mut key = keys::data_key(new_prefix);
-        let new_prefix_data_key_len = key.len();
+        let mut data_key = keys::DATA_PREFIX_KEY.to_vec();
+        let data_key_prefix_len = keys::DATA_PREFIX_KEY.len();
+        let mut user_key = new_prefix.to_vec();
+        let user_key_prefix_len = new_prefix.len();
         let mut first_key = None;
 
         match range_start {
@@ -1016,23 +1041,44 @@ impl SstImporter {
             .unwrap();
 
         while iter.valid()? {
-            let old_key = keys::origin_key(iter.key());
-            if is_after_end_bound(old_key, &range_end) {
+            let mut old_key = Cow::Borrowed(keys::origin_key(iter.key()));
+            let mut ts = None;
+
+            if is_after_end_bound(old_key.as_ref(), &range_end) {
                 break;
             }
+
+            if req_type == DownloadRequestType::Keyspace {
+                ts = Some(Key::decode_ts_bytes_from(old_key.as_ref())?.to_owned());
+                old_key = {
+                    let mut key = old_key.to_vec();
+                    decode_bytes_in_place(&mut key, false)?;
+                    Cow::Owned(key)
+                };
+            }
+
             if !old_key.starts_with(old_prefix) {
                 return Err(Error::WrongKeyPrefix {
                     what: "Key in SST",
                     key: keys::origin_key(iter.key()).to_vec(),
                     prefix: old_prefix.to_vec(),
                 });
             }
-            key.truncate(new_prefix_data_key_len);
-            key.extend_from_slice(&old_key[old_prefix.len()..]);
+
+            data_key.truncate(data_key_prefix_len);
+            user_key.truncate(user_key_prefix_len);
+            user_key.extend_from_slice(&old_key[old_prefix.len()..]);
+            if req_type == DownloadRequestType::Keyspace {
+                data_key.extend(encode_bytes(&user_key));
+                data_key.extend(ts.unwrap());
+            } else {
+                data_key.extend_from_slice(&user_key);
+            }
+
             let mut value = Cow::Borrowed(iter.value());
 
             if rewrite_rule.new_timestamp != 0 {
-                key = Key::from_encoded(key)
+                data_key = Key::from_encoded(data_key)
                     .truncate_ts()
                     .map_err(|e| {
                         Error::BadFormat(format!(
@@ -1056,10 +1102,10 @@ impl SstImporter {
                 }
             }
 
-            sst_writer.put(&key, &value)?;
+            sst_writer.put(&data_key, &value)?;
             iter.next()?;
             if first_key.is_none() {
-                first_key = Some(keys::origin_key(&key).to_vec());
+                first_key = Some(keys::origin_key(&data_key).to_vec());
             }
         }
 
@@ -1078,7 +1124,7 @@ impl SstImporter {
 
             let mut final_range = Range::default();
             final_range.set_start(start_key);
-            final_range.set_end(keys::origin_key(&key).to_vec());
+            final_range.set_end(keys::origin_key(&data_key).to_vec());
             Ok(Some(final_range))
         } else {
             // nothing is written: prevents finishing the SST at all.

components/sst_importer/src/sst_importer.rs

@@ -192,6 +192,16 @@ impl Key {
         Ok(number::decode_u64_desc(&mut ts)?.into())
     }
 
+    /// Decode the timestamp from a ts encoded key and return in bytes.
+    #[inline]
+    pub fn decode_ts_bytes_from(key: &[u8]) -> Result<&[u8], codec::Error> {
+        let len = key.len();
+        if len < number::U64_SIZE {
+            return Err(codec::Error::KeyLength);
+        }
+        Ok(&key[key.len() - number::U64_SIZE..])
+    }
+
     /// Whether the user key part of a ts encoded key `ts_encoded_key` equals to
     /// the encoded user key `user_key`.
     ///

components/txn_types/src/types.rs

@@ -783,7 +783,9 @@ where
                 cipher,
                 limiter,
                 engine,
-                DownloadExt::default().cache_key(req.get_storage_cache_id()),
+                DownloadExt::default()
+                    .cache_key(req.get_storage_cache_id())
+                    .req_type(req.get_request_type()),
             );
             let mut resp = DownloadResponse::default();
             match res.await {