// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0. // Modifications Copyright 2022 PingCAP, Inc. use std::{borrow::Cow, marker::PhantomData}; use bytes::{Buf, Bytes}; use kvengine::{Item, LOCK_CF, SnapAccess, UserMeta, ValueCacheValue, WRITE_CF, read}; use kvproto::kvrpcpb::IsolationLevel; use tikv_kv::{Snapshot, Statistics}; use tikv_util::{Either, txn_debug}; use txn_types::{ Key, Lock, LockType, OldValue, SHORT_VALUE_MAX_LEN, TimeStamp, TsSet, Value, ValueEntry, Write, WriteType, is_short_value, }; use crate::storage::{ REQUEST_EXCEED_BOUND, config::ReadRescheduleOptions, mvcc, mvcc::NewerTsCheckState, txn::{Error, ErrorInner, Result, TxnEntry, TxnEntryStore}, }; pub struct CloudStore { marker: PhantomData, snapshot: kvengine::SnapAccess, value_cache: Option, start_ts: u64, bypass_locks: TsSet, fill_cache: bool, stats: Statistics, incremental_get_met_newer_ts_data: NewerTsCheckState, check_has_newer_ts_data: bool, read_reschedule_options: ReadRescheduleOptions, } #[maybe_async::async_trait] impl super::Store for CloudStore { type Scanner = CloudStoreScanner; #[maybe_async] async fn get(&self, user_key: &Key, statistics: &mut Statistics) -> Result> { txn_debug!( trace_event::types::Category::ReadDetails, "CloudStore::get entry"; "key" => log_wrappers::Value::key(user_key.as_encoded()), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id() ); // No need check newer ts data for normal get. let mut disable_newer_ts_check = NewerTsCheckState::Unknown; let item = Self::get_inner( user_key, &self.snapshot, self.start_ts, &self.bypass_locks, statistics, self.value_cache.as_ref(), &mut disable_newer_ts_check, ) .await?; let result = if item.value_len() > 0 { Some(item.get_value().to_vec()) } else { None }; txn_debug!( trace_event::types::Category::ReadDetails, "CloudStore::get result"; "key" => log_wrappers::Value::key(user_key.as_encoded()), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id(), "result" => ?result.as_ref().map(|v| log_wrappers::Value::value(v)) ); Ok(result) } #[maybe_async] async fn incremental_get(&mut self, user_key: &Key) -> Result> { txn_debug!( trace_event::types::Category::ReadDetails, "CloudStore::incremental_get entry"; "key" => log_wrappers::Value::key(user_key.as_encoded()), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id() ); let stat = &mut self.stats; let item = Self::get_inner( user_key, &self.snapshot, self.start_ts, &self.bypass_locks, stat, self.value_cache.as_ref(), &mut self.incremental_get_met_newer_ts_data, ) .await?; let result = if item.value_len() > 0 { Some(item.get_value().to_vec()) } else { None }; txn_debug!(trace_event::types::Category::ReadDetails, "CloudStore::incremental_get result"; "key" => log_wrappers::Value::key(user_key.as_encoded()), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id(), "result" => ?result.as_ref().map(|v| log_wrappers::Value::value(v)) ); Ok(result) } #[maybe_async] async fn incremental_get_entry( &mut self, user_key: &Key, load_commit_ts: bool, ) -> Result> { txn_debug!( trace_event::types::Category::ReadDetails, "CloudStore::incremental_get_entry entry"; "key" => log_wrappers::Value::key(user_key.as_encoded()), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id(), "load_commit_ts" => load_commit_ts ); let item = Self::get_inner( user_key, &self.snapshot, self.start_ts, &self.bypass_locks, &mut self.stats, self.value_cache.as_ref(), &mut self.incremental_get_met_newer_ts_data, ) .await?; let result = if item.value_len() > 0 { let commit_ts = if load_commit_ts { if item.user_meta_len() == 0 { return Err(ErrorInner::Other(box_err!( "CloudStore visible value is missing user metadata" )) .into()); } Some(TimeStamp::new( UserMeta::from_slice(item.user_meta()).commit_ts, )) } else { None }; Some(ValueEntry::new(item.get_value().to_vec(), commit_ts)) } else { None }; txn_debug!( trace_event::types::Category::ReadDetails, "CloudStore::incremental_get_entry result"; "key" => log_wrappers::Value::key(user_key.as_encoded()), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id(), "result" => ?result .as_ref() .map(|entry| log_wrappers::Value::value(&entry.value)), "commit_ts" => ?result.as_ref().and_then(|entry| entry.commit_ts) ); Ok(result) } fn incremental_get_take_statistics(&mut self) -> Statistics { std::mem::take(&mut self.stats) } fn incremental_get_statistics(&self) -> Statistics { self.stats.clone() } fn incremental_get_met_newer_ts_data(&self) -> NewerTsCheckState { self.incremental_get_met_newer_ts_data } async fn batch_get( &self, keys: &[Key], statistics: &mut Vec, ) -> Result>>> { self.trace_batch_get_entry(keys); let res_vec = super::store::batch_get_with_reschedule( self, keys, statistics, self.read_reschedule_options .batch_get_fairness_checkpoint_keys, ) .await?; self.trace_batch_get_result(keys, &res_vec); Ok(res_vec) } #[maybe_async] async fn scanner( &mut self, desc: bool, _key_only: bool, check_has_newer_ts_data: bool, load_commit_ts: bool, lower_bound: Option, upper_bound: Option, ) -> Result { txn_debug!(trace_event::types::Category::ReadDetails, "CloudStore::scanner entry"; "desc" => desc, "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id(), "lower_bound" => ?lower_bound, "upper_bound" => ?upper_bound ); self.scanner_inner( desc, lower_bound, upper_bound, false, check_has_newer_ts_data, load_commit_ts, ) .await } fn get_kvengine_snap(&self) -> Option { Some(self.snapshot.clone()) } fn is_snapshot_sync(&self) -> bool { self.snapshot.is_sync() } fn get_read_ts(&self) -> u64 { self.start_ts } #[inline] fn is_check_has_newer_ts_data(&self) -> bool { self.check_has_newer_ts_data } } impl CloudStore { /// Creates a CloudStore over a region snapshot. /// /// `check_has_newer_ts_data` must only be enabled for coprocessor requests /// whose response may be returned as cacheable. A checked incremental /// point get starts from `NotMetYet` and switches to `Met` if the read /// observes MVCC state newer than `start_ts`. pub fn new( snapshot: S, start_ts: u64, bypass_locks: TsSet, fill_cache: bool, check_has_newer_ts_data: bool, ) -> Self { Self { marker: PhantomData, snapshot: snapshot.get_kvengine_snap().unwrap().clone(), value_cache: snapshot.get_value_cache().cloned(), start_ts, bypass_locks, fill_cache, stats: Statistics::default(), incremental_get_met_newer_ts_data: if check_has_newer_ts_data { NewerTsCheckState::NotMetYet } else { NewerTsCheckState::Unknown }, check_has_newer_ts_data, read_reschedule_options: ReadRescheduleOptions::default(), } } /// Applies a request-stable snapshot of the online read reschedule /// settings. pub fn with_read_reschedule_options(mut self, options: ReadRescheduleOptions) -> Self { self.read_reschedule_options = options; self } fn trace_batch_get_entry(&self, keys: &[Key]) { txn_debug!(trace_event::types::Category::ReadDetails, "CloudStore::batch_get entry"; "keys_count" => keys.len(), "keys" => ?keys, "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id() ); } fn trace_batch_get_result(&self, keys: &[Key], results: &[Result>]) { txn_debug!(trace_event::types::Category::ReadDetails, "CloudStore::batch_get result"; "keys_count" => keys.len(), "start_ts" => ?self.start_ts, "keyspace_id" => self.snapshot.get_keyspace_id(), "region_id" => self.snapshot.get_id(), "result" => ?results.iter().map(|r| { match r { Ok(Some(value)) => Ok(log_wrappers::Value::value(value)), Ok(None) => Ok(log_wrappers::Value::value(b"")), Err(e) => Err(e) } }).collect::>() ); } #[maybe_async::both] async fn get_inner<'a>( user_key: &Key, snap: &'a kvengine::SnapAccess, start_ts: u64, bypass_locks: &TsSet, statistics: &mut Statistics, value_cache: Option<&'a kvengine::ValueCache>, met_newer_ts_data: &mut NewerTsCheckState, ) -> mvcc::Result> { tikv_util::set_current_region(snap.get_id()); let raw_key = user_key.to_raw()?; let item = snap.get( LOCK_CF, &raw_key, snap.get_mem_table_snap_version().into_inner(), ); statistics.lock.get += 1; statistics.lock.flow_stats.read_keys += 1; statistics.lock.flow_stats.read_bytes += raw_key.len() + item.value_len(); statistics.lock.processed_keys += 1; let has_lock = item.value_len() > 0; if has_lock { let lock_or_shared_locks = txn_types::parse_lock(item.get_value())?; // A lock makes the response timestamp-dependent even when the // current read timestamp can bypass it. TiDB's coprocessor cache // key intentionally ignores the read timestamp, so a later request // with the same payload could observe the lock after it commits or // resolves. Marking it uncacheable preserves the data-version // contract: unchanged snapshot version is enough only when no // newer MVCC state was hidden by this read timestamp. mark_met_newer_ts_data(met_newer_ts_data); txn_types::check_ts_conflict( Cow::Borrowed(&lock_or_shared_locks), user_key, TimeStamp::new(start_ts), bypass_locks, IsolationLevel::Si, )?; } if snap.get_start_key() > raw_key.as_slice() || snap.get_end_key() <= raw_key.as_slice() { error!( "get key {:?} out of snap range {:?}, {:?}, {}:{}", raw_key.as_slice(), snap.get_start_key(), snap.get_end_key(), snap.get_id(), snap.get_version() ); return Ok(Item::default()); } if *met_newer_ts_data == NewerTsCheckState::NotMetYet { // Point get normally seeks at `start_ts`, which hides versions with // larger commit_ts. Before returning the value visible at // `start_ts`, peek for any newer write on the same raw key. Newer // puts and deletes both mean a response for this timestamp cannot // be reused by a later timestamp just because the snapshot data // version still matches. `has_newer` intentionally does not resolve // a hidden blob value because only version existence matters here. if snap.has_newer(WRITE_CF, &raw_key, start_ts).await { mark_met_newer_ts_data(met_newer_ts_data); } } let item = if let Some(cache) = value_cache { if let Some(val) = cache.get(snap, &raw_key, start_ts, has_lock) { let hit_item = if cache.is_double_check() { let origin_item = snap.get(WRITE_CF, &raw_key, start_ts).await; if !cached_value_is_logically_equivalent(&origin_item, &val) { error!( "value cache mismatch: read missing or invalid"; "key" => log_wrappers::Value::key(&raw_key), "region_id" => snap.get_id(), "shard_ver" => snap.get_version(), "keyspace_id" => snap.get_keyspace_id(), "start_ts" => start_ts, "cached_version" => val.version, "read_version" => origin_item.version, ); debug_assert!(false); } origin_item } else { Item::from_cached_value(&val) }; return Ok(hit_item); } // We use the u64::MAX to get first to fill the cache, if version is newer, // use start_ts to get again. If the we use start_ts first to get the value, // we do not know if the value can be cached, because the value may be updated // after the start_ts. let mut item = snap.get(WRITE_CF, &raw_key, u64::MAX).await; if item.version > start_ts { item = snap.get(WRITE_CF, &raw_key, start_ts).await; } else if item.user_meta_len() > 0 && !has_lock { let um = UserMeta::from_slice(item.user_meta()); let val = ValueCacheValue::new( Bytes::copy_from_slice(item.get_value()), um.start_ts, item.version, snap.get_mem_table_snap_version(), ); cache.set(snap, &raw_key, val); } item } else { snap.get(WRITE_CF, &raw_key, start_ts).await }; statistics.write.get += 1; statistics.write.flow_stats.read_keys += 1; statistics.write.flow_stats.read_bytes += user_key.len() + item.value_len(); if item.value_len() > 0 { statistics.write.processed_keys += 1; statistics.processed_size += user_key.len() + item.value_len(); } statistics.total_size += user_key.len() + item.value_len(); Ok(item) } pub fn check_locks_in_range( &self, lower_bound: &[u8], upper_bound: &[u8], check_new_ts: bool, ) -> Result<()> { let lower_bound = Some(Key::from_raw(lower_bound)); let upper_bound = Some(Key::from_raw(upper_bound)); let (lower_bound, upper_bound) = verify_range(&self.snapshot, lower_bound, upper_bound)?; let mut stats = Statistics::default(); let mut iter = self .snapshot .new_iterator(LOCK_CF, false, false, None, self.fill_cache); if iter.set_range(lower_bound, upper_bound) { stats.lock.seek += 1; } stats.lock.seek_tombstone += iter.take_skipped_tombstones(); // A defense in depth, asserts that `check_locks_in_range` is not // called with `check_new_ts` enabled. As `check_locks_in_range` is only // used by `delegate_coprocessor` and the status server's dump-snapshot // API. Neither uses the coprocessor cache, checking newer-ts data // is unnecessary here. assert!( !check_new_ts, "check_locks_in_range should not check newer ts data" ); let mut disable_newer_ts_check = NewerTsCheckState::Unknown; check_locks( &mut iter, &self.bypass_locks, self.start_ts, &mut stats, &mut disable_newer_ts_check, )?; Ok(()) } pub fn new_memtable_iterator(&self) -> kvengine::read::Iterator { self.snapshot .new_memtable_iterator(WRITE_CF, false, false, Some(self.start_ts)) } #[maybe_async::both] async fn scanner_inner( &self, desc: bool, lower_bound: Option, upper_bound: Option, output_delete: bool, check_has_newer_ts_data: bool, load_commit_ts: bool, ) -> Result { tikv_util::set_current_region(self.snapshot.get_id()); let scanner = CloudStoreScanner::new( self.snapshot.clone(), desc, self.fill_cache, self.bypass_locks.clone(), self.start_ts, lower_bound, upper_bound, output_delete, check_has_newer_ts_data, load_commit_ts, ) .await?; Ok(scanner.with_read_reschedule_options(self.read_reschedule_options)) } #[cfg(test)] pub fn get_start_ts(&self) -> u64 { self.start_ts } #[cfg(test)] pub fn get_by_pass_locks(&self) -> TsSet { self.bypass_locks.clone() } #[cfg(test)] pub fn is_fill_cache(&self) -> bool { self.fill_cache } } /// Checks lock conflicts for the iterator's current range. /// /// When `met_newer_ts_data` is provided, any observed lock is also recorded as /// newer MVCC state for coprocessor cache validation. The response can depend /// on whether the read timestamp bypasses that lock, while the coprocessor /// cache key intentionally does not include the timestamp. pub fn check_locks( lock_iter: &mut read::Iterator, bypass_locks: &TsSet, start_ts: u64, stats: &mut Statistics, met_newer_ts_data: &mut NewerTsCheckState, ) -> mvcc::Result<()> { while lock_iter.valid() { check_current_lock(lock_iter, bypass_locks, start_ts, stats, met_newer_ts_data)?; lock_iter.next(); stats.lock.next_tombstone += lock_iter.take_skipped_tombstones(); } Ok(()) } fn mark_met_newer_ts_data(met_newer_ts_data: &mut NewerTsCheckState) { if *met_newer_ts_data == NewerTsCheckState::NotMetYet { *met_newer_ts_data = NewerTsCheckState::Met; } } fn check_current_lock( lock_iter: &mut read::Iterator, bypass_locks: &TsSet, start_ts: u64, stats: &mut Statistics, met_newer_ts_data: &mut NewerTsCheckState, ) -> mvcc::Result<()> { let raw_key = lock_iter.key(); let key = Key::from_raw(raw_key); let raw_key_len = raw_key.len(); let val = lock_iter.val(); stats.lock.flow_stats.read_keys += 1; stats.lock.flow_stats.read_bytes += raw_key_len + val.len(); stats.lock.processed_keys += 1; let lock_or_shared_locks = txn_types::parse_lock(val)?; // Scans check the whole requested lock range before returning rows. Any // observed lock is newer MVCC state that the response may depend on, so // it disables coprocessor caching even if `check_ts_conflict` lets this // specific read timestamp continue. mark_met_newer_ts_data(met_newer_ts_data); txn_types::check_ts_conflict( Cow::Borrowed(&lock_or_shared_locks), &key, start_ts.into(), bypass_locks, IsolationLevel::Si, )?; Ok(()) } fn verify_range( snap: &SnapAccess, lower_bound: Option, upper_bound: Option, ) -> Result<(Bytes, Bytes)> { let lower = if let Some(lower_key) = &lower_bound { let lower = Bytes::from(lower_key.to_raw()?); if lower.chunk() < snap.get_start_key() { range_error(snap, &lower_bound, &upper_bound)?; } lower } else { snap.get_start_key().to_vec().into() }; let upper = if let Some(upper_key) = &upper_bound { let upper = Bytes::from(upper_key.to_raw()?); if upper.chunk() > snap.get_end_key() { range_error(snap, &lower_bound, &upper_bound)?; } upper } else { snap.get_end_key().to_vec().into() }; Ok((lower, upper)) } fn range_error( snap: &SnapAccess, lower_bound: &Option, upper_bound: &Option, ) -> Result<()> { REQUEST_EXCEED_BOUND.inc(); let start = lower_bound.as_ref().map(|k| k.as_encoded().clone()); let end = upper_bound.as_ref().map(|k| k.as_encoded().clone()); let snap_start = Key::from_raw(snap.get_start_key()); let snap_end = Key::from_raw(snap.get_end_key()); Err(Error::from(ErrorInner::InvalidReqRange { start, end, lower_bound: Some(snap_start.into_encoded()), upper_bound: Some(snap_end.into_encoded()), })) } fn cached_value_is_logically_equivalent( origin_item: &Item<'_>, cached_value: &ValueCacheValue, ) -> bool { // Version zero means the authoritative read found no physical item. An empty // cached tombstone has the same logical value, so only its bytes must match. (origin_item.version == 0 || origin_item.version == cached_value.version) && origin_item.get_value() == cached_value.value.chunk() } // DeltaScanner scans a key range and yields all MVCC versions of each key /// after (but excluded) `start_ts` with representing: /// - committed versions (from Write CF) as `TxnEntry::Commit` (including /// deletes), /// - uncommitted prewrite locks (from Lock CF) as `TxnEntry::Prewrite`. pub struct CloudDeltaScanner { lock_iter: kvengine::read::Iterator, iter: kvengine::read::Iterator, start_ts: u64, pub stats: Statistics, lower_bound: Bytes, upper_bound: Bytes, filter_by_ts: bool, is_started: bool, current_raw_key: Vec, snap: SnapAccess, } pub enum DeltaEntry { Version { user_meta: UserMeta, value: Vec, encoded_key: Key, old_value: OldValue, }, Lock { lock: Lock, raw_key: Vec, old_value: OldValue, }, } impl std::fmt::Debug for DeltaEntry { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Version { user_meta, value, encoded_key, old_value, } => f .debug_struct("Version") .field("user_meta", user_meta) .field("value", &format_args!("<{} byte(s)>", value.len())) .field("encoded_key", encoded_key) .field("old_value", old_value) .finish(), Self::Lock { lock, raw_key, old_value, } => f .debug_struct("Lock") .field("lock", lock) .field("raw_key", &log_wrappers::Value::key(raw_key)) .field("old_value", old_value) .finish(), } } } impl DeltaEntry { fn into_txn_entry(self) -> TxnEntry { match self { DeltaEntry::Version { user_meta, value, encoded_key, old_value, } => { let write_key = encoded_key.clone().append_ts(user_meta.commit_ts.into()); let write_ty = if value.is_empty() { WriteType::Delete } else { WriteType::Put }; let default_key = encoded_key.append_ts(user_meta.start_ts.into()); if value.len() < SHORT_VALUE_MAX_LEN { let write_record = Write::new(write_ty, user_meta.start_ts.into(), Some(value)); TxnEntry::Commit { default: (vec![], vec![]), write: (write_key.into_encoded(), write_record.as_ref().to_bytes()), old_value, } } else { let write_record = Write::new(write_ty, user_meta.start_ts.into(), None); let default_value = value; TxnEntry::Commit { default: (default_key.into_encoded(), default_value), write: (write_key.into_encoded(), write_record.as_ref().to_bytes()), old_value, } } } DeltaEntry::Lock { mut lock, raw_key, old_value, } => { assert!( lock.short_value.is_some() || lock.lock_type == LockType::Delete, "lock value in cloud engine without short value; lock = {:?}, key = {}", lock, log_wrappers::Value::key(&raw_key) ); let encoded_key = Key::from_raw(&raw_key).into_encoded(); if lock.short_value.as_ref().map(|v| v.len()).unwrap_or(0) > SHORT_VALUE_MAX_LEN { let val = lock.short_value.take().unwrap(); let default_key = Key::from_raw(&raw_key).append_ts(lock.ts); return TxnEntry::Prewrite { default: (default_key.into_encoded(), val), lock: (encoded_key, lock.to_bytes()), old_value, }; } TxnEntry::Prewrite { default: (vec![], vec![]), lock: (encoded_key, lock.to_bytes()), old_value, } } } } } impl CloudDeltaScanner { #[maybe_async::both] pub async fn new( snap: SnapAccess, start_ts: u64, lower_bound: Option, upper_bound: Option, filter_by_ts: bool, ) -> Self { let stats = Statistics::default(); let lock_iter = snap.new_iterator(LOCK_CF, false, false, None, false); let iter = if filter_by_ts { snap.new_delta_iterator(false, start_ts) } else { snap.new_iterator(WRITE_CF, false, true, None, false) }; let (lower_bound, upper_bound) = verify_range(&snap, lower_bound, upper_bound).unwrap_or((Bytes::new(), Bytes::new())); Self { lock_iter, start_ts, iter, lower_bound, upper_bound, stats, filter_by_ts, is_started: false, current_raw_key: vec![], snap, } } #[maybe_async::both] pub async fn init(&mut self) -> Result<()> { if self .lock_iter .set_range(self.lower_bound.clone(), self.upper_bound.clone()) { self.stats.lock.seek += 1; } self.stats.lock.seek_tombstone += self.lock_iter.take_skipped_tombstones(); if self .iter .set_range(self.lower_bound.clone(), self.upper_bound.clone()) { self.stats.write.seek += 1; } self.is_started = true; if self.iter.valid() { self.current_raw_key = self.iter.key().to_vec(); } Ok(()) } #[maybe_async::both] async fn get_old_value(&self, key: &[u8], before_ts: u64) -> OldValue { let item = self.snap.get(WRITE_CF, key, before_ts - 1).await; let v = item.get_value(); if v.is_empty() { OldValue::None } else { OldValue::value(v.to_vec()) } } #[maybe_async::both] pub async fn next_inner(&mut self) -> Result> { if !self.is_started { self.init().await?; } loop { // As we assume the lock cf's kv count should be relatively small, we use // `get_old_value` instead of reusing the write cf's iterator to // simplify the code. if self.lock_iter.valid() { let raw_key_len = self.lock_iter.key().len(); let val_len = self.lock_iter.val().len(); if self.current_raw_key.is_empty() // Write iter reachs the end... // or it exceeds the lock iter. (We should yield Prewrite records before Commits.) || self.lock_iter.key() <= self.current_raw_key.as_slice() { let lock = match txn_types::parse_lock(self.lock_iter.val()) .map_err(Error::from_mvcc)? { Either::Left(lock) => lock, Either::Right(_) => { self.lock_iter.next().await; continue; } }; // ignore pessimistic or SELECT FOR UPDATE lock. if lock.lock_type == LockType::Pessimistic || lock.lock_type == LockType::Lock { self.lock_iter.next().await; self.stats.lock.next_tombstone += self.lock_iter.take_skipped_tombstones(); self.stats.lock.next += 1; self.stats.lock.flow_stats.read_keys += 1; self.stats.lock.flow_stats.read_bytes += raw_key_len + val_len; continue; } let lock_ts = std::cmp::max(lock.ts, lock.for_update_ts); let entry = DeltaEntry::Lock { lock, raw_key: self.lock_iter.key().to_vec(), // NOTE: maybe read it from the `WRITE_CF` state. // if `lock.key == write.key` => old value is `write.val` old_value: self .get_old_value(self.lock_iter.key(), lock_ts.into_inner()) .await, }; self.lock_iter.next().await; self.stats.lock.next_tombstone += self.lock_iter.take_skipped_tombstones(); self.stats.lock.next += 1; self.stats.lock.flow_stats.read_keys += 1; self.stats.lock.flow_stats.read_bytes += raw_key_len + val_len; self.stats.lock.processed_keys += 1; return Ok(Some(entry)); } } if self.iter.valid() { let iter_key = self.iter.key(); let key_len = iter_key.len(); let key = Key::from_raw(iter_key); let user_meta = UserMeta::from_slice(self.iter.user_meta()); if user_meta.commit_ts <= self.start_ts { debug!("delta scanner get a commit record less than start_ts"; "key" => log_wrappers::Value::key(key.as_encoded()), "user_meta" => ?user_meta, "start_ts" => ?self.start_ts, ); self.iter.next().await; if !self.iter.valid() || self.iter.key() != self.current_raw_key { self.current_raw_key.truncate(0); if self.iter.valid() { self.current_raw_key.extend_from_slice(self.iter.key()); } } continue; } let val = self.iter.val(); let val_len = val.len(); let val_vec = val.to_vec(); self.iter.next().await; self.stats.write.next += 1; self.stats.write.flow_stats.read_keys += 1; self.stats.write.flow_stats.read_bytes += key_len + val_len; let old_value = if !self.iter.valid() || self.iter.key() != self.current_raw_key { let old_value = if self.filter_by_ts { self.get_old_value(&self.current_raw_key, user_meta.commit_ts) .await } else { OldValue::None }; self.stats.write.processed_keys += 1; self.current_raw_key.truncate(0); if self.iter.valid() { self.current_raw_key.extend_from_slice(self.iter.key()); } old_value } else { // Here we are in the previous version. let prev = self.iter.val(); if prev.is_empty() { OldValue::None } else { OldValue::value(prev.to_vec()) } }; self.stats.processed_size += key_len + val_len; let entry = DeltaEntry::Version { user_meta, value: val_vec, encoded_key: key, old_value, }; return Ok(Some(entry)); } return Ok(None); } } } impl super::TxnEntryScanner for CloudDeltaScanner { fn next_entry(&mut self) -> Result> { Ok(self.next_inner()?.map(DeltaEntry::into_txn_entry)) } fn take_statistics(&mut self) -> Statistics { std::mem::take(&mut self.stats) } } pub struct CloudStoreScanner { snap: SnapAccess, lock_iter: kvengine::read::Iterator, bypass_locks: TsSet, start_ts: u64, iter: kvengine::read::Iterator, stats: Statistics, is_started: bool, lower_bound: Bytes, upper_bound: Bytes, output_delete: bool, met_newer_ts_data: NewerTsCheckState, load_commit_ts: bool, fair_cpu_checkpoint_entries: usize, } impl CloudStoreScanner { #[maybe_async::both] pub async fn new( snap: SnapAccess, desc: bool, fill_cache: bool, bypass_locks: TsSet, start_ts: u64, lower_bound: Option, upper_bound: Option, output_delete: bool, check_has_newer_ts_data: bool, load_commit_ts: bool, ) -> Result { let stats = Statistics::default(); let lock_iter = snap.new_iterator(LOCK_CF, desc, false, None, fill_cache); let iter = snap .new_iterator(WRITE_CF, desc, false, Some(start_ts), fill_cache) .await; let (lower_bound, upper_bound) = verify_range(&snap, lower_bound, upper_bound)?; Ok(Self { snap, lock_iter, bypass_locks, start_ts, iter, stats, is_started: false, lower_bound, upper_bound, output_delete, met_newer_ts_data: if check_has_newer_ts_data { NewerTsCheckState::NotMetYet } else { NewerTsCheckState::Unknown }, load_commit_ts, fair_cpu_checkpoint_entries: ReadRescheduleOptions::default() .cloud_store_fair_cpu_checkpoint_entries, }) } fn with_read_reschedule_options(mut self, options: ReadRescheduleOptions) -> Self { self.fair_cpu_checkpoint_entries = options.cloud_store_fair_cpu_checkpoint_entries; self } #[maybe_async::both] async fn init(&mut self, checkpointed_entries: &mut usize) -> Result<()> { if self .lock_iter .set_range(self.lower_bound.clone(), self.upper_bound.clone()) { self.stats.lock.seek += 1; }; self.stats.lock.seek_tombstone += self.lock_iter.take_skipped_tombstones(); if self.take_reschedule_checkpoint_due(checkpointed_entries) { self.check_reschedule().await; } while self.lock_iter.valid() { check_current_lock( &mut self.lock_iter, &self.bypass_locks, self.start_ts, &mut self.stats, &mut self.met_newer_ts_data, )?; self.lock_iter.next(); self.stats.lock.next_tombstone += self.lock_iter.take_skipped_tombstones(); if self.take_reschedule_checkpoint_due(checkpointed_entries) { self.check_reschedule().await; } } if self .iter .set_range(self.lower_bound.clone(), self.upper_bound.clone()) .await { self.stats.write.seek += 1; } if self.take_reschedule_checkpoint_due(checkpointed_entries) { self.check_reschedule().await; } Ok(()) } #[inline] fn scanned_entries(&self) -> usize { self.lock_iter .iter_stats() .keys .saturating_add(self.iter.iter_stats().keys) } fn take_reschedule_checkpoint_due(&self, checkpointed_entries: &mut usize) -> bool { let scanned_entries = self.scanned_entries(); if self.fair_cpu_checkpoint_entries == 0 || scanned_entries.saturating_sub(*checkpointed_entries) < self.fair_cpu_checkpoint_entries { return false; } *checkpointed_entries = scanned_entries; true } #[inline] fn check_reschedule(&mut self) {} async fn check_reschedule_async(&mut self) { tikv_util::cooperative_checkpoint().await; } #[maybe_async::both] async fn next_inner(&mut self) -> Result> { let mut checkpointed_entries = self.scanned_entries(); if self.is_started { self.iter.next().await; if self.take_reschedule_checkpoint_due(&mut checkpointed_entries) { self.check_reschedule().await; } } else { self.init(&mut checkpointed_entries).await?; self.is_started = true; } loop { if self.met_newer_ts_data == NewerTsCheckState::NotMetYet && self.iter.met_newer_ts() { self.met_newer_ts_data = NewerTsCheckState::Met; } if !self.iter.valid() { return Ok(None); } let iter_key = self.iter.key(); let key_len = iter_key.len(); let key = Key::from_raw(iter_key); let user_meta = UserMeta::from_slice(self.iter.user_meta()); let val = self.iter.val(); self.stats.write.flow_stats.read_keys += 1; self.stats.write.flow_stats.read_bytes += key_len + val.len(); if !val.is_empty() || self.output_delete { // Tick processed_key because it's visible to user. self.stats.write.processed_keys += 1; self.stats.processed_size += key_len + val.len(); return Ok(Some((key, user_meta, val.to_vec()))); } self.stats.write.next_tombstone += 1; self.iter.next().await; if self.take_reschedule_checkpoint_due(&mut checkpointed_entries) { self.check_reschedule().await; } } } pub fn next_with_user_meta(&mut self) -> Result> { self.next_inner() } } #[maybe_async::async_trait] impl super::Scanner for CloudStoreScanner { #[maybe_async] async fn next(&mut self) -> Result> { let result = self .next_inner() .await? .map(|(key, _user_meta, val)| (key, val)); txn_debug!(trace_event::types::Category::ReadDetails, "CloudStoreScanner::next"; "start_ts" => ?self.start_ts, "keyspace_id" => self.snap.get_keyspace_id(), "region_id" => self.snap.get_id(), "result" => ?result.as_ref().map(|(key, value)| { (key, log_wrappers::Value::value(value)) }) ); Ok(result) } #[maybe_async] async fn next_entry(&mut self) -> Result> { let result = self.next_inner().await?.map(|(key, user_meta, value)| { let commit_ts = self .load_commit_ts .then(|| TimeStamp::new(user_meta.commit_ts)); (key, ValueEntry::new(value, commit_ts)) }); txn_debug!(trace_event::types::Category::ReadDetails, "CloudStoreScanner::next_entry"; "start_ts" => ?self.start_ts, "keyspace_id" => self.snap.get_keyspace_id(), "region_id" => self.snap.get_id(), "result" => ?result.as_ref().map(|(key, entry)| { ( key, log_wrappers::Value::value(&entry.value), entry.commit_ts, ) }) ); Ok(result) } fn is_snapshot_sync(&self) -> bool { self.snap.is_sync() } fn met_newer_ts_data(&self) -> NewerTsCheckState { self.met_newer_ts_data } fn take_statistics(&mut self) -> Statistics { let mut stats = std::mem::take(&mut self.stats); let lock_stats = self.lock_iter.take_iter_stats(); let write_stats = self.iter.take_iter_stats(); stats.total_size += lock_stats.total_size + write_stats.total_size; stats.lock.next = lock_stats.keys; stats.write.next = write_stats.keys; stats } fn statistics(&self) -> Statistics { let mut stats = self.stats.clone(); stats.total_size += self.lock_iter.iter_stats().total_size + self.iter.iter_stats().total_size; stats } fn reset_range( &mut self, desc: bool, load_commit_ts: bool, lower_bound: Option, upper_bound: Option, ) -> Result { if self.iter.is_reverse() != desc { return Ok(false); } let (lower_bound, upper_bound) = verify_range(&self.snap, lower_bound, upper_bound)?; self.lower_bound = lower_bound; self.upper_bound = upper_bound; self.load_commit_ts = load_commit_ts; self.is_started = false; Ok(true) } } impl super::TxnEntryScanner for CloudStoreScanner { fn next_entry(&mut self) -> Result> { Ok(self.next_inner()?.map(|(key, user_meta, val)| { let write_key = key.append_ts(TimeStamp::new(user_meta.commit_ts)); let write_type = if val.is_empty() { WriteType::Delete } else { WriteType::Put }; let (write, default_key, default_val) = if is_short_value(&val) { ( Write::new(write_type, user_meta.start_ts.into(), Some(val)), vec![], vec![], ) } else { let default_key = write_key .clone() .truncate_ts() .unwrap() .append_ts(user_meta.start_ts.into()); ( Write::new(write_type, user_meta.start_ts.into(), None), default_key.into_encoded(), val, ) }; TxnEntry::Commit { default: (default_key, default_val), write: (write_key.into_encoded(), write.as_ref().to_bytes()), old_value: OldValue::None, } })) } fn take_statistics(&mut self) -> Statistics { std::mem::take(&mut self.stats) } } impl TxnEntryStore for CloudStore { type Scanner = CloudStoreScanner; fn entry_scanner( &self, lower_bound: Option, upper_bound: Option, _after_ts: TimeStamp, output_delete: bool, ) -> Result { self.scanner_inner(false, lower_bound, upper_bound, output_delete, false, false) } } #[cfg(test)] mod tests { use std::{ sync::{ Arc, atomic::{AtomicUsize, Ordering}, }, time::Duration, }; use cpu_time::ThreadTime; use futures::executor::block_on; use kvengine::{ WriteBatch, test_engine::{ ApplyTask, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX, TestEngine, install_lock_cf_test_table, new_test_engine_opt, }, }; use rfstore::store::RegionSnapshot; use tikv_util::{ mpsc, yatp_pool::{ DefaultTicker, FairAdmitOutcome, FairCpuAdmissionMetadata, FairCpuBudgetAction, FairCpuPermit, FairQueueAdmissionController, FairQueueAdmissionResult, TaskMetadata, YatpPoolBuilder, set_task_meta, }, }; use txn_types::{LockType, TimeStamp}; use yatp::{queue::Extras, task::future::TaskCell}; use super::*; use crate::storage::txn::{ Scanner as TxnScanner, Store as TxnStore, TxnEntryScanner, nextgen_scanned_bytes, }; #[test] fn test_nextgen_scanned_bytes_uses_total_versions_size() { let mut stats = Statistics { processed_size: 10, total_size: 30, ..Default::default() }; assert_eq!(nextgen_scanned_bytes(&stats), 30); stats.processed_size = 40; assert_eq!(nextgen_scanned_bytes(&stats), 40); } #[test] fn test_cached_value_logical_equivalence() { struct Case { name: &'static str, origin: Option<(u64, &'static [u8])>, cached_version: u64, cached_value: &'static [u8], equivalent: bool, } let cases = [ Case { name: "missing origin matches empty cached value", origin: None, cached_version: 2, cached_value: b"", equivalent: true, }, Case { name: "missing origin differs from nonempty cached value", origin: None, cached_version: 2, cached_value: b"value", equivalent: false, }, Case { name: "matching version and value", origin: Some((2, b"value")), cached_version: 2, cached_value: b"value", equivalent: true, }, Case { name: "different nonzero version", origin: Some((2, b"value")), cached_version: 3, cached_value: b"value", equivalent: false, }, Case { name: "different value", origin: Some((2, b"value")), cached_version: 2, cached_value: b"other", equivalent: false, }, ]; for case in cases { let origin_item = match case.origin { Some((version, value)) => Item::from_cached_value(&ValueCacheValue::new( Bytes::from_static(value), 1, version, Default::default(), )), None => Item::default(), }; let cached_value = ValueCacheValue::new( Bytes::from_static(case.cached_value), 1, case.cached_version, Default::default(), ); assert_eq!( cached_value_is_logically_equivalent(&origin_item, &cached_value), case.equivalent, "{}", case.name ); } } fn lock_bytes_with_type(primary: &[u8], start_ts: u64, lock_type: LockType) -> Vec { let short_value = match lock_type { LockType::Put => Some(b"value".to_vec()), LockType::Delete | LockType::Lock | LockType::Pessimistic | LockType::Shared => None, }; Lock::new( lock_type, primary.to_vec(), TimeStamp::new(start_ts), 3000, short_value, TimeStamp::zero(), 1, TimeStamp::zero(), ) .to_bytes() } fn lock_bytes(primary: &[u8], start_ts: u64) -> Vec { lock_bytes_with_type(primary, start_ts, LockType::Put) } fn write_data(wb: WriteBatch, applier_tx: &mpsc::Sender) { let (result_tx, result_rx) = mpsc::bounded(1); applier_tx.send(ApplyTask::new_wb(wb, result_tx)).unwrap(); result_rx.recv().unwrap().unwrap(); } fn switch_mem_table(applier_tx: &mpsc::Sender) { let mut wb = WriteBatch::default(); wb.reset(1); wb.set_switch_mem_table(); write_data(wb, applier_tx); } fn write_value( applier_tx: &mpsc::Sender, key: &[u8], value: &[u8], start_ts: u64, commit_ts: u64, ) { let mut wb = WriteBatch::default(); wb.reset(1); wb.put( WRITE_CF, key, value, 0, &UserMeta::new(start_ts, commit_ts).to_array(), commit_ts, ); write_data(wb, applier_tx); } fn write_delete( applier_tx: &mpsc::Sender, key: &[u8], start_ts: u64, commit_ts: u64, ) { let mut wb = WriteBatch::default(); wb.reset(1); wb.put( WRITE_CF, key, &[], 0, &UserMeta::new(start_ts, commit_ts).to_array(), commit_ts, ); write_data(wb, applier_tx); } fn write_lock(applier_tx: &mpsc::Sender, key: &[u8], start_ts: u64) { let mut wb = WriteBatch::default(); wb.reset(1); wb.put( LOCK_CF, key, &lock_bytes_with_type(key, start_ts, LockType::Lock), 0, &[], 0, ); write_data(wb, applier_tx); } fn cloud_store_with_cache_check( engine: &TestEngine, start_ts: u64, check_has_newer_ts_data: bool, ) -> CloudStore { let snap = engine.get_shard(1).unwrap().new_snap_access(); let region_snapshot = RegionSnapshot::from_snapshot(snap, None, None); CloudStore::new( region_snapshot, start_ts, TsSet::Empty, false, check_has_newer_ts_data, ) } fn cloud_store(engine: &TestEngine, start_ts: u64) -> CloudStore { cloud_store_with_cache_check(engine, start_ts, true) } fn scanner_with_cache_check( engine: &TestEngine, start_ts: u64, lower: usize, upper: usize, check_has_newer_ts_data: bool, ) -> CloudStoreScanner { let kb = engine.key_builder(); let snap = engine.get_shard(1).unwrap().new_snap_access(); CloudStoreScanner::new( snap, false, false, TsSet::Empty, start_ts, Some(Key::from_raw(&kb.i_to_outer_key(lower))), Some(Key::from_raw(&kb.i_to_outer_key(upper))), false, check_has_newer_ts_data, false, ) .unwrap() } #[test] fn test_cloud_store_enables_newer_ts_cacheability_check() { let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let store = cloud_store(&engine, 200); assert!(store.is_check_has_newer_ts_data()); } #[test] fn test_incremental_get_cacheability_tracks_newer_writes() { let (engine, applier_tx) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let key = engine.key_builder().i_to_outer_key(10); write_value(&applier_tx, &key, b"v1", 100, 200); write_value(&applier_tx, &key, b"v2", 300, 400); // An old read can return v1, but it must not be cached because a later // timestamp with the same request payload would return v2 while the // request cache key ignores the timestamp. let mut old_store = cloud_store(&engine, 200); assert_eq!( old_store.incremental_get(&Key::from_raw(&key)).unwrap(), Some(b"v1".to_vec()) ); assert_eq!( old_store.incremental_get_met_newer_ts_data(), NewerTsCheckState::Met ); // Reading at the latest committed version has no hidden newer MVCC // state, so it remains cacheable if the snapshot data version matches. let mut latest_store = cloud_store(&engine, 400); assert_eq!( latest_store.incremental_get(&Key::from_raw(&key)).unwrap(), Some(b"v2".to_vec()) ); assert_eq!( latest_store.incremental_get_met_newer_ts_data(), NewerTsCheckState::NotMetYet ); // When cache checking is disabled, the state is deliberately unknown // and the response layer will not mark the result cacheable. let mut unchecked_store = cloud_store_with_cache_check(&engine, 200, false); assert_eq!( unchecked_store .incremental_get(&Key::from_raw(&key)) .unwrap(), Some(b"v1".to_vec()) ); assert_eq!( unchecked_store.incremental_get_met_newer_ts_data(), NewerTsCheckState::Unknown ); } #[test] fn test_incremental_get_cacheability_tracks_newer_delete_and_lock() { let (engine, applier_tx) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let deleted_key = engine.key_builder().i_to_outer_key(20); write_value(&applier_tx, &deleted_key, b"v1", 100, 200); switch_mem_table(&applier_tx); write_delete(&applier_tx, &deleted_key, 300, 400); let mut delete_store = cloud_store(&engine, 200); let _ = delete_store .incremental_get(&Key::from_raw(&deleted_key)) .unwrap(); assert_eq!( delete_store.incremental_get_met_newer_ts_data(), NewerTsCheckState::Met ); let locked_key = engine.key_builder().i_to_outer_key(30); write_value(&applier_tx, &locked_key, b"v1", 100, 200); write_lock(&applier_tx, &locked_key, 500); let mut lock_store = cloud_store(&engine, 200); assert_eq!( lock_store .incremental_get(&Key::from_raw(&locked_key)) .unwrap(), Some(b"v1".to_vec()) ); assert_eq!( lock_store.incremental_get_met_newer_ts_data(), NewerTsCheckState::Met ); } #[test] fn test_scan_cacheability_tracks_newer_writes_and_deletes() { let (engine, applier_tx) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let key1 = kb.i_to_outer_key(40); let key2 = kb.i_to_outer_key(41); write_value(&applier_tx, &key1, b"v1", 100, 200); write_value(&applier_tx, &key2, b"v2", 300, 400); // key2 is invisible at start_ts=200, but its newer committed version is // inside the scanned range, so the partial result is not reusable by a // later timestamp. let mut scanner = scanner_with_cache_check(&engine, 200, 40, 42, true); assert_eq!( scanner .next() .unwrap() .map(|(k, v)| (k.into_raw().unwrap(), v)), Some((key1.clone(), b"v1".to_vec())) ); assert!(scanner.next().unwrap().is_none()); assert_eq!(scanner.met_newer_ts_data(), NewerTsCheckState::Met); let key3 = kb.i_to_outer_key(50); write_value(&applier_tx, &key3, b"v3", 100, 200); switch_mem_table(&applier_tx); write_delete(&applier_tx, &key3, 300, 400); let mut delete_scanner = scanner_with_cache_check(&engine, 200, 50, 51, true); let _ = delete_scanner.next().unwrap(); assert_eq!(delete_scanner.met_newer_ts_data(), NewerTsCheckState::Met); let mut latest_scanner = scanner_with_cache_check(&engine, 400, 40, 42, true); assert!(latest_scanner.next().unwrap().is_some()); assert_eq!( latest_scanner.met_newer_ts_data(), NewerTsCheckState::NotMetYet ); } #[test] fn test_scan_cacheability_tracks_locks_and_disabled_check() { let (engine, applier_tx) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let key1 = kb.i_to_outer_key(60); let locked_key = kb.i_to_outer_key(61); write_value(&applier_tx, &key1, b"v1", 100, 200); write_lock(&applier_tx, &locked_key, 500); let mut scanner = scanner_with_cache_check(&engine, 200, 60, 62, true); assert!(scanner.next().unwrap().is_some()); assert_eq!(scanner.met_newer_ts_data(), NewerTsCheckState::Met); let key2 = kb.i_to_outer_key(70); write_value(&applier_tx, &key2, b"v2", 100, 200); write_value(&applier_tx, &key2, b"v3", 300, 400); let mut unchecked_scanner = scanner_with_cache_check(&engine, 200, 70, 71, false); assert!(unchecked_scanner.next().unwrap().is_some()); assert_eq!( unchecked_scanner.met_newer_ts_data(), NewerTsCheckState::Unknown ); } struct YieldOnCheckpoint; impl FairCpuPermit for YieldOnCheckpoint { fn consume(&mut self, _cpu_delta_us: u64) -> FairCpuBudgetAction { FairCpuBudgetAction::YieldToScheduler } } #[derive(Default)] struct YieldingAdmissionController { admissions: AtomicUsize, } impl FairQueueAdmissionController for YieldingAdmissionController { fn precheck(&self) -> FairQueueAdmissionResult { FairQueueAdmissionResult::Ready } fn admit( &self, _keyspace_id: u32, _task_type: u8, _deadline: Option, ) -> FairAdmitOutcome { self.admissions.fetch_add(1, Ordering::SeqCst); FairAdmitOutcome::Admitted(Box::new(YieldOnCheckpoint)) } } #[test] fn test_lock_scan_yields() { const LOCK_COUNT: usize = 256; const CHECKPOINT_ENTRIES: usize = 32; let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let first_key_index = 300; let locks = (0..LOCK_COUNT) .map(|offset| { let key_index = first_key_index + offset; let key = kb.i_to_outer_key(key_index); (key_index, lock_bytes(&key, 10)) }) .collect::>(); install_lock_cf_test_table(&engine, 1001, &locks); let snap = engine.get_shard(1).unwrap().new_snap_access(); let mut scanner = block_on(CloudStoreScanner::new_async( snap, false, false, TsSet::All, 1, Some(Key::from_raw(&kb.i_to_outer_key(first_key_index))), Some(Key::from_raw( &kb.i_to_outer_key(first_key_index + LOCK_COUNT), )), false, false, false, )) .unwrap(); scanner.fair_cpu_checkpoint_entries = CHECKPOINT_ENTRIES; let controller = Arc::new(YieldingAdmissionController::default()); let mut builder = YatpPoolBuilder::new(DefaultTicker); builder .name_prefix("test-cloud-lock-scan-yields") .thread_count(1, 1, 1); let (pool, _queue) = builder.build_keyspace_fair_pool(Some(controller.clone())); let (tx, rx) = mpsc::unbounded(); let future = async move { // Make the CPU-time rate limit due before the first entry-based // checkpoint. { let timer = ThreadTime::now(); while timer.elapsed() < Duration::from_millis(2) { std::hint::black_box(()); } } let result = scanner.next_async().await; let stats = TxnScanner::take_statistics(&mut scanner); tx.send((result, stats)).unwrap(); }; let mut extras = Extras::new_multilevel(1, None); set_task_meta( &mut extras, &TaskMetadata { keyspace_id: 1, fair_cpu_admission: Some(FairCpuAdmissionMetadata { task_type: 1, deadline: None, }), }, ); pool.spawn(TaskCell::new(future, extras)); let (result, stats) = rx.recv_timeout(Duration::from_secs(10)).unwrap(); assert!(result.unwrap().is_none()); assert_eq!(stats.lock.next, LOCK_COUNT); assert_eq!(stats.lock.processed_keys, LOCK_COUNT); assert!( controller.admissions.load(Ordering::SeqCst) > 1, "the lock loop should yield and be admitted again before it finishes" ); } #[test] fn test_lock_cf_tombstone_statistics_count() { let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let active_key = kb.i_to_outer_key(100); let next_active_key = kb.i_to_outer_key(105); // Keep the fixture deterministic: one LOCK_CF SST, no L0/level overlap, // and exactly four tombstones between two visible locks. install_lock_cf_test_table( &engine, 1000, &[ (100, lock_bytes(&active_key, 10)), (101, vec![]), (102, vec![]), (103, vec![]), (104, vec![]), (105, lock_bytes(&next_active_key, 11)), ], ); let snap = engine.get_shard(1).unwrap().new_snap_access(); // Seeking into the tombstone-only span should land on the next visible // lock and charge all skipped tombstones to seek_tombstone. let mut scanner = CloudDeltaScanner::new( snap.clone(), 0, Some(Key::from_raw(&kb.i_to_outer_key(101))), Some(Key::from_raw(&kb.i_to_outer_key(106))), true, ); scanner.init().unwrap(); let stats = scanner.take_statistics(); assert_eq!(stats.lock.seek, 1); assert_eq!(stats.lock.seek_tombstone, 4); assert_eq!(stats.lock.next_tombstone, 0); // Starting from a visible lock exercises the next() path. After the // first lock is emitted, advancing the lock iterator crosses the same // four tombstones, so they should be charged to next_tombstone. let mut scanner = CloudDeltaScanner::new( snap, 0, Some(Key::from_raw(&active_key)), Some(Key::from_raw(&kb.i_to_outer_key(106))), true, ); scanner.init().unwrap(); let stats = scanner.take_statistics(); assert_eq!(stats.lock.seek, 1); assert_eq!(stats.lock.seek_tombstone, 0); assert_eq!(stats.lock.next, 0); assert_eq!(stats.lock.next_tombstone, 0); assert!(matches!( scanner.next_entry().unwrap(), Some(TxnEntry::Prewrite { .. }) )); let stats = scanner.take_statistics(); assert_eq!(stats.lock.seek_tombstone, 0); assert_eq!(stats.lock.next, 1); assert_eq!(stats.lock.next_tombstone, 4); // Bottom-level checkpoint accounting peeks at the same non-draining // iterator statistics used by final request statistics, including // internally skipped tombstones. let snap = engine.get_shard(1).unwrap().new_snap_access(); let mut scanner = CloudStoreScanner::new( snap, false, false, TsSet::All, u64::MAX, Some(Key::from_raw(&active_key)), Some(Key::from_raw(&kb.i_to_outer_key(106))), false, false, false, ) .unwrap(); assert!(scanner.next().unwrap().is_none()); assert_eq!(scanner.scanned_entries(), 6); let stats = TxnScanner::take_statistics(&mut scanner); assert_eq!(stats.lock.next + stats.write.next, 6); } #[test] fn test_lock_cf_filtered_lock_statistics_count() { let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let filtered_key = kb.i_to_outer_key(200); let active_key = kb.i_to_outer_key(203); let filtered_lock = lock_bytes_with_type(&filtered_key, 20, LockType::Lock); let active_lock = lock_bytes(&active_key, 21); let expected_read_bytes = filtered_key.len() + filtered_lock.len() + active_key.len() + active_lock.len(); // The first visible entry is a lock-only prewrite. It should be read // and counted in CfStatistics, but filtered out from scanner output. // The following two tombstones make the skipped next() path observable. install_lock_cf_test_table( &engine, 1000, &[ (200, filtered_lock), (201, vec![]), (202, vec![]), (203, active_lock), ], ); let snap = engine.get_shard(1).unwrap().new_snap_access(); let mut scanner = CloudDeltaScanner::new( snap, 0, Some(Key::from_raw(&filtered_key)), Some(Key::from_raw(&kb.i_to_outer_key(204))), true, ); scanner.init().unwrap(); let stats = scanner.take_statistics(); assert_eq!(stats.lock.seek, 1); assert_eq!(stats.lock.seek_tombstone, 0); assert_eq!(stats.lock.next, 0); assert_eq!(stats.lock.next_tombstone, 0); assert!(matches!( scanner.next_entry().unwrap(), Some(TxnEntry::Prewrite { .. }) )); let stats = scanner.take_statistics(); assert_eq!(stats.lock.next, 2); assert_eq!(stats.lock.next_tombstone, 2); assert_eq!(stats.lock.flow_stats.read_keys, 2); assert_eq!(stats.lock.flow_stats.read_bytes, expected_read_bytes); assert_eq!(stats.lock.processed_keys, 1); } }