// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0. // Modifications Copyright 2022 PingCAP, Inc. use std::{ops::Deref, sync::Arc}; use bytes::Bytes; use kvengine::{EXTRA_CF, LOCK_CF, UserMeta, WRITE_CF, table::InnerKey}; use raftstore::store::TxnExt; use tikv_util::Either; use txn_types::{Key, Lock, LockOrSharedLocks, OldValue, TimeStamp, Value, Write, WriteType}; use crate::storage::{ config::{ExtraCfScanMode, ReadRescheduleOptions}, kv::Error as KvError, mvcc::{ ErrorInner, Result, TxnCommitRecord, metrics::{ EXTRA_CF_CONFLICT_CHECK_COUNTER, EXTRA_CF_PRUNING_DATASETS_COUNTER, EXTRA_CF_SCAN_EXAMINED_ENTRIES, ExtraCfConflictOutcomeKind, ExtraCfScanModeKind, ExtraCfScanStrategyKind, }, }, }; // Context for scanning in-memory pessimistic locks. #[derive(Clone)] pub struct InMemoryLockContext { pub txn_ext: Arc, pub term: u64, pub version: u64, } /// Cooperative safe point for the lock scan. /// /// The pair is hand-written rather than generated: `maybe_async` rewrites every /// awaited call by appending `_async`, so the checkpoint has to be reached /// through a name that already follows that convention. Taking no arguments /// keeps `CloudReader`, which is not `Sync`, out of the generated future. #[inline] fn check_lock_scan_reschedule() {} async fn check_lock_scan_reschedule_async() { tikv_util::cooperative_checkpoint().await; } fn validate_extra_key_matches_raw_key(raw_key: &[u8], extra_key: &[u8]) -> Result<()> { if extra_key.len() == raw_key.len() + 8 { return Ok(()); } Err(ErrorInner::Other(box_err!( "key prefix violation implies possible interleaving in EXTRA CF: expected raw_key({} bytes) + timestamp(8 bytes) = {} bytes, found {} bytes, raw_key: {}, found key: {}", raw_key.len(), raw_key.len() + 8, extra_key.len(), log_wrappers::Value::key(raw_key), log_wrappers::Value::value(extra_key) )) .into()) } /// Parameters for EXTRA_CF conflict scanning. struct ScanParams { /// Exclusive lower bound for dataset-level pruning. pruning_threshold: TimeStamp, /// Exclusive lower bound for write-conflict detection. conflict_threshold: TimeStamp, /// Detect self-rollback inline when set. detect_self_rollback: Option, } pub struct CloudReader { pub(crate) snapshot: kvengine::SnapAccess, fill_cache: bool, extra_cf_scan_mode: ExtraCfScanMode, pub statistics: tikv_kv::Statistics, /// Cached EXTRA_CF iterator for reuse across multiple legacy conflict /// checks. cached_extra_cf_iter: Option, /// Number of LOCK_CF iterator entries between cooperative checkpoints in /// the lock scan. Zero disables the checkpoints. fair_cpu_checkpoint_entries: usize, } impl CloudReader { pub fn new(snapshot: kvengine::SnapAccess, fill_cache: bool) -> Self { Self::new_with_scan_mode(snapshot, fill_cache, ExtraCfScanMode::default()) } pub fn new_with_scan_mode( snapshot: kvengine::SnapAccess, fill_cache: bool, extra_cf_scan_mode: ExtraCfScanMode, ) -> Self { Self { snapshot, fill_cache, extra_cf_scan_mode, statistics: tikv_kv::Statistics::default(), cached_extra_cf_iter: None, fair_cpu_checkpoint_entries: ReadRescheduleOptions::default() .cloud_store_fair_cpu_checkpoint_entries, } } /// Applies the caller's snapshot of the read reschedule intervals. /// /// Callers that run outside a read-pool poll can skip this: the /// cooperative checkpoint is a no-op without a poll-scoped context. pub fn with_read_reschedule_options(mut self, options: ReadRescheduleOptions) -> Self { self.fair_cpu_checkpoint_entries = options.cloud_store_fair_cpu_checkpoint_entries; self } /// Returns whether the lock scan has consumed one checkpoint interval /// worth of LOCK_CF entries since the previous checkpoint. /// /// The cadence follows the iterator's entry counter rather than the number /// of returned locks: one `Iterator::next` can walk an unbounded number of /// physical entries while skipping tombstones and filtered-out locks, so /// counting loop iterations would place no checkpoint at all on a range /// that returns few locks. fn take_lock_scan_checkpoint_due( &self, lock_iter: &kvengine::read::Iterator, checkpointed_entries: &mut usize, ) -> bool { if self.fair_cpu_checkpoint_entries == 0 { return false; } let scanned_entries = lock_iter.iter_stats().keys; if scanned_entries.saturating_sub(*checkpointed_entries) < self.fair_cpu_checkpoint_entries { return false; } *checkpointed_entries = scanned_entries; true } #[inline] fn extra_cf_mode_label(mode: ExtraCfScanMode) -> ExtraCfScanModeKind { match mode { ExtraCfScanMode::Legacy => ExtraCfScanModeKind::legacy, ExtraCfScanMode::Pruned => ExtraCfScanModeKind::pruned, ExtraCfScanMode::Both => ExtraCfScanModeKind::both, } } #[inline] fn extra_cf_strategy_label(strategy: ExtraCfScanStrategyKind) -> &'static str { match strategy { ExtraCfScanStrategyKind::unified => "unified", ExtraCfScanStrategyKind::two_step => "two_step", } } fn compare_both_mode_results( strategy: ExtraCfScanStrategyKind, full: Option, pruned: Option, ) -> Option { let strategy = Self::extra_cf_strategy_label(strategy); match (full, pruned) { (None, None) => None, (Some(full_um), Some(pruned_um)) if full_um.is_rollback() && pruned_um.is_rollback() => { if full_um.start_ts != pruned_um.start_ts { panic!( "EXTRA_CF both-mode mismatch ({strategy}): legacy rollback start_ts {} != pruned rollback start_ts {}", full_um.start_ts, pruned_um.start_ts, ); } Some(pruned_um) } (Some(full_um), Some(pruned_um)) if full_um.is_rollback() != pruned_um.is_rollback() => { panic!( "EXTRA_CF both-mode mismatch ({strategy}): legacy rollback flag {} != pruned rollback flag {}; legacy={:?}, pruned={:?}", full_um.is_rollback(), pruned_um.is_rollback(), full_um, pruned_um, ); } (Some(full_um), Some(pruned_um)) => { if full_um.start_ts != pruned_um.start_ts || full_um.commit_ts != pruned_um.commit_ts { panic!( "EXTRA_CF both-mode mismatch ({strategy}): legacy write conflict {:?} != pruned write conflict {:?}", full_um, pruned_um, ); } Some(pruned_um) } (full, pruned) => { panic!( "EXTRA_CF both-mode mismatch ({strategy}): legacy={:?}, pruned={:?}", full, pruned, ); } } } fn run_extra_cf_conflict_scan( &mut self, raw_key: &[u8], params: &ScanParams, strategy: ExtraCfScanStrategyKind, ) -> Result> { match self.extra_cf_scan_mode { ExtraCfScanMode::Legacy => self.scan_extra_cf_for_write_conflict_full(raw_key, params), ExtraCfScanMode::Pruned => { self.scan_extra_cf_for_write_conflict_pruned(raw_key, params) } ExtraCfScanMode::Both => { let full = self.scan_extra_cf_for_write_conflict_full(raw_key, params)?; let pruned = self.scan_extra_cf_for_write_conflict_pruned(raw_key, params)?; Ok(Self::compare_both_mode_results(strategy, full, pruned)) } } } fn observe_extra_cf_pruned_iter_stats(stats: &kvengine::read::TsFilterStats) { Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.candidate.memtable, stats.candidate_memtables, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.candidate.l0, stats.candidate_l0s, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.candidate.stable, stats.candidate_stables, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.filtered.memtable, stats.pruned_memtables_by_max_ts, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.filtered.l0, stats.pruned_l0s_by_max_ts, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.filtered.stable, stats.pruned_stables_by_max_ts, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.selected.memtable, stats.surviving_memtables, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.selected.l0, stats.surviving_l0s, ); Self::inc_extra_cf_pruning_datasets( &EXTRA_CF_PRUNING_DATASETS_COUNTER.selected.stable, stats.surviving_stables, ); } fn inc_extra_cf_pruning_datasets(counter: &prometheus::IntCounter, count: usize) { if count == 0 { return; } counter.inc_by(count as u64); } #[inline] fn observe_extra_cf_scan_examined_entries(mode: ExtraCfScanMode, entries: usize) { EXTRA_CF_SCAN_EXAMINED_ENTRIES .get(Self::extra_cf_mode_label(mode)) .observe(entries as f64); } #[inline] fn observe_extra_cf_conflict_check( mode: ExtraCfScanMode, strategy: ExtraCfScanStrategyKind, outcome: ExtraCfConflictOutcomeKind, ) { EXTRA_CF_CONFLICT_CHECK_COUNTER .get(Self::extra_cf_mode_label(mode)) .get(strategy) .get(outcome) .inc(); } fn get_commit_by_item( user_meta: &UserMeta, value: &[u8], start_ts: TimeStamp, ) -> Option { if user_meta.start_ts == start_ts.into_inner() { let (write_type, short_value) = if value.is_empty() { (WriteType::Delete, None) } else { (WriteType::Put, Some(value.to_vec())) }; let write = Write::new(write_type, TimeStamp::new(user_meta.start_ts), short_value); return Some(TxnCommitRecord::SingleRecord { commit_ts: TimeStamp::new(user_meta.commit_ts), write, }); } None } /// Note: This method is also used by resolving locks during restoring /// keyspace. #[maybe_async::both] pub async fn get_txn_commit_record( &mut self, key: &Key, start_ts: TimeStamp, ) -> Result { let raw_key = key.to_raw()?; let item = self.snapshot.get(WRITE_CF, &raw_key, 0).await; if item.user_meta_len() > 0 { let user_meta = UserMeta::from_slice(item.user_meta()); if let Some(record) = Self::get_commit_by_item(&user_meta, item.get_value(), start_ts) { return Ok(record); } } let mut data_iter = self .snapshot .new_iterator(WRITE_CF, false, true, None, self.fill_cache) .await; let mut next_key = Vec::with_capacity(raw_key.len() + 1); next_key.extend_from_slice(&raw_key); next_key.push(0); data_iter .set_range(Bytes::from(raw_key), Bytes::from(next_key)) .await; while data_iter.valid() { debug_assert!(!kvengine::table::is_deleted(data_iter.meta())); // TODO: remove this check, iterator should not return deleted records. if kvengine::table::is_deleted(data_iter.meta()) { break; } let user_meta = UserMeta::from_slice(data_iter.user_meta()); if user_meta.commit_ts < start_ts.into_inner() { // A transaction's commit_ts must be greater than start_ts, if current commit_ts // is already smaller than the start_ts, we don't need to look for older // version. break; } if let Some(record) = Self::get_commit_by_item(&user_meta, data_iter.val(), start_ts) { return Ok(record); } data_iter.next().await; } match self.get_extra(key, start_ts) { Some((commit_ts, write)) => Ok(TxnCommitRecord::SingleRecord { commit_ts, write }), None => Ok(TxnCommitRecord::None { overlapped_write: None, }), } } /// If the physical time gap between `max_allowed_write` and `start_ts` /// exceeds this threshold, fall back to the two-step approach (point GET + /// scan with `max_allowed_write` as threshold) instead of the unified scan /// (scan with `start_ts - 1` as threshold). /// /// Rationale: the unified scan uses a lower pruning threshold (`start_ts - /// 1`) which retains more datasets. When the gap is large (e.g. pessimistic /// transactions with many for_update_ts bumps), the extra datasets can make /// the scan significantly slower, outweighing the saved point GET. const UNIFIED_SCAN_MAX_PHYSICAL_GAP_MS: u64 = 1000; /// Check for write conflicts in EXTRA_CF. /// Returns the first conflict found (either self-rollback or newer write). /// Self-rollback is not a global priority over newer conflicts: in the /// unified strategy it is detected inline and follows EXTRA_CF iteration /// order. The two-step strategy performs the exact self-rollback lookup /// before scanning for newer conflicts. /// /// `Legacy` restores the pre-pruning conflict-selection behavior by always /// using one ordered full-source scan. `Pruned` and `Both` use a heuristic /// to choose between two strategies: /// - **Unified scan** (when `max_allowed_write` is close to `start_ts`): /// single scan with conservative pruning, detecting self-rollback inline. /// Saves one point GET in the common no-conflict path. /// - **Two-step** (when `max_allowed_write >> start_ts`): separate point /// GET for self-rollback, then scan with `threshold = max_allowed_write` /// for better pruning. /// /// # Parameters /// - `start_ts`: The start_ts of the transaction to be checked for /// self-rollback /// - `max_allowed_write`: Exclusive threshold - any write with commit_ts > /// this value is a conflict pub fn find_extra_cf_conflict_record( &mut self, key: &Key, start_ts: TimeStamp, max_allowed_write: TimeStamp, ) -> Result> { let raw_key = key.to_raw()?; self.find_extra_cf_conflict_record_with_raw_key(&raw_key, start_ts, max_allowed_write) } pub fn find_extra_cf_conflict_record_with_raw_key( &mut self, raw_key: &[u8], start_ts: TimeStamp, max_allowed_write: TimeStamp, ) -> Result> { if self.extra_cf_scan_mode == ExtraCfScanMode::Legacy { return self.find_extra_cf_conflict_unified(raw_key, start_ts, max_allowed_write); } let physical_gap = max_allowed_write .physical() .saturating_sub(start_ts.physical()); if physical_gap <= Self::UNIFIED_SCAN_MAX_PHYSICAL_GAP_MS { self.find_extra_cf_conflict_unified(raw_key, start_ts, max_allowed_write) } else { self.find_extra_cf_conflict_two_step(raw_key, start_ts, max_allowed_write) } } /// Two-step strategy: point GET for self-rollback, then scan with /// `max_allowed_write` as the pruning/conflict threshold. fn find_extra_cf_conflict_two_step( &mut self, raw_key: &[u8], start_ts: TimeStamp, max_allowed_write: TimeStamp, ) -> Result> { if let Some(user_meta) = self.check_self_rollback(raw_key, start_ts)? { Self::observe_extra_cf_conflict_check( self.extra_cf_scan_mode, ExtraCfScanStrategyKind::two_step, ExtraCfConflictOutcomeKind::self_rollback, ); return Ok(Some(user_meta)); } let params = ScanParams { pruning_threshold: max_allowed_write, conflict_threshold: max_allowed_write, detect_self_rollback: None, }; let result = self.run_extra_cf_conflict_scan(raw_key, ¶ms, ExtraCfScanStrategyKind::two_step)?; Self::observe_extra_cf_conflict_check( self.extra_cf_scan_mode, ExtraCfScanStrategyKind::two_step, if result.is_some() { ExtraCfConflictOutcomeKind::write_conflict } else { ExtraCfConflictOutcomeKind::no_conflict }, ); Ok(result) } /// Unified strategy: scan from `start_ts.prev()` to include self-rollback. fn find_extra_cf_conflict_unified( &mut self, raw_key: &[u8], start_ts: TimeStamp, max_allowed_write: TimeStamp, ) -> Result> { let pruning_threshold = if max_allowed_write < start_ts { warn!( "unexpected EXTRA_CF conflict check threshold below start_ts"; "start_ts" => start_ts, "max_allowed_write" => max_allowed_write, ); max_allowed_write } else { start_ts.prev() }; let params = ScanParams { pruning_threshold, conflict_threshold: max_allowed_write, detect_self_rollback: Some(start_ts), }; let result = self.run_extra_cf_conflict_scan(raw_key, ¶ms, ExtraCfScanStrategyKind::unified)?; let outcome = match &result { Some(um) if um.is_rollback() => ExtraCfConflictOutcomeKind::self_rollback, Some(_) => ExtraCfConflictOutcomeKind::write_conflict, None => ExtraCfConflictOutcomeKind::no_conflict, }; Self::observe_extra_cf_conflict_check( self.extra_cf_scan_mode, ExtraCfScanStrategyKind::unified, outcome, ); Ok(result) } // Backward-compatible point-get for EXTRA_CF. // // EXTRA_CF encoding changed in PR #3484 (reverse-bits -> TiKV-style desc). // To avoid missing historical records, try the alternative encoding when // the primary lookup misses. fn lookup_extra_user_meta(&self, raw_key: &[u8], start_ts: TimeStamp) -> Option { let start_ts_u64 = start_ts.into_inner(); let extra_key = kvengine::encode_extra_txn_status_key(raw_key, start_ts_u64); let item = self.snapshot.get(EXTRA_CF, &extra_key, 0); if item.exists() { return Some(UserMeta::from_slice(item.user_meta())); } let candidates = [ kvengine::encode_extra_txn_status_key_revbits(raw_key, start_ts_u64), kvengine::encode_extra_txn_status_key_desc(raw_key, start_ts_u64), ]; for candidate in candidates { if candidate.as_ref() == extra_key.as_ref() { continue; } let item = self.snapshot.get(EXTRA_CF, candidate.as_ref(), 0); if item.exists() { return Some(UserMeta::from_slice(item.user_meta())); } } None } /// Check for self-rollback using point lookup optimization. fn check_self_rollback( &mut self, raw_key: &[u8], start_ts: TimeStamp, ) -> Result> { Ok(self .lookup_extra_user_meta(raw_key, start_ts) .filter(UserMeta::is_rollback)) } /// Legacy implementation: scan EXTRA_CF via the normal high-level snapshot /// iterator path, i.e. the same path used when `extra_cf_scan_mode=legacy`. fn scan_extra_cf_for_write_conflict_full( &mut self, raw_key: &[u8], params: &ScanParams, ) -> Result> { let conflict_threshold = params.conflict_threshold.into_inner(); let mut extra_iter = self.cached_extra_cf_iter.take().unwrap_or_else(|| { self.snapshot .new_iterator(EXTRA_CF, false, false, None, self.fill_cache) }); extra_iter.seek(raw_key); let mut iterations = 0usize; let mut result = None; // Raw keys can interleave in the EXTRA CF. // This loop handles it well by checking the prefix and the length match. // NOTE: raw keys in EXTRA CF must not interleave. This depends on the // property that all TiDB keys are mem-comparable, i.e. their order won't change // when appending a ts. In other words, no key can be a prefix of // another key. while extra_iter.valid() && extra_iter.key().starts_with(raw_key) { iterations += 1; validate_extra_key_matches_raw_key(raw_key, extra_iter.key())?; if kvengine::table::is_deleted(extra_iter.meta()) { extra_iter.next(); continue; } let um = UserMeta::from_slice(extra_iter.user_meta()); if let Some(self_start_ts) = params.detect_self_rollback { if um.is_rollback() && um.start_ts == self_start_ts.into_inner() { result = Some(um); break; } } if !um.is_rollback() && um.commit_ts > conflict_threshold { result = Some(um); break; } extra_iter.next(); } Self::observe_extra_cf_scan_examined_entries(ExtraCfScanMode::Legacy, iterations); self.cached_extra_cf_iter = Some(extra_iter); Ok(result) } /// Scan EXTRA_CF for write conflicts using pruned iterator sources. /// Returns immediately upon finding the first conflict. fn scan_extra_cf_for_write_conflict_pruned( &mut self, raw_key: &[u8], params: &ScanParams, ) -> Result> { let raw_inner_key = InnerKey::from_outer_key(raw_key); let conflict_threshold = params.conflict_threshold.into_inner(); let mut prune_stats = kvengine::read::TsFilterStats::default(); let mut extra_iter = self.snapshot.build_pruned_extra_cf_iter( raw_key, params.pruning_threshold.into_inner(), self.fill_cache, &mut prune_stats, ); Self::observe_extra_cf_pruned_iter_stats(&prune_stats); extra_iter.seek(raw_inner_key); if !extra_iter.valid() { Self::observe_extra_cf_scan_examined_entries(ExtraCfScanMode::Pruned, 0); return Ok(None); } let mut iterations = 0usize; // NOTE: raw keys in EXTRA CF must not interleave. This depends on the // property that all TiDB keys are mem-comparable, i.e. their order // won't change when appending a ts. In other words, no key can be a // prefix of another key. while extra_iter.valid() && extra_iter.key().starts_with(raw_inner_key.deref()) { iterations += 1; validate_extra_key_matches_raw_key(raw_inner_key.deref(), extra_iter.key().deref())?; if extra_iter.value().is_deleted() { extra_iter.next(); continue; } let um = UserMeta::from_slice(extra_iter.value().user_meta()); if let Some(self_start_ts) = params.detect_self_rollback { if um.is_rollback() && um.start_ts == self_start_ts.into_inner() { Self::observe_extra_cf_scan_examined_entries( ExtraCfScanMode::Pruned, iterations, ); return Ok(Some(um)); } } if !um.is_rollback() && um.commit_ts > conflict_threshold { Self::observe_extra_cf_scan_examined_entries(ExtraCfScanMode::Pruned, iterations); return Ok(Some(um)); } extra_iter.next(); } Self::observe_extra_cf_scan_examined_entries(ExtraCfScanMode::Pruned, iterations); Ok(None) } // do not use it for conflict check, use `check_write_conflict_in_extra_cf` // instead. pub fn get_extra(&mut self, key: &Key, start_ts: TimeStamp) -> Option<(TimeStamp, Write)> { let raw_key = key.to_raw().unwrap(); self.get_extra_with_raw_key(&raw_key, start_ts) } pub fn get_extra_with_raw_key( &mut self, raw_key: &[u8], start_ts: TimeStamp, ) -> Option<(TimeStamp, Write)> { let user_meta = self.lookup_extra_user_meta(raw_key, start_ts)?; Some(if user_meta.is_rollback() { (start_ts, Write::new(WriteType::Rollback, start_ts, None)) } else { ( user_meta.commit_ts.into(), Write::new(WriteType::Lock, start_ts, None), ) }) } pub fn load_lock(&mut self, key: &Key) -> Result> { let raw_key = key.to_raw().unwrap(); self.load_lock_with_raw_key(&raw_key) } pub fn load_lock_with_raw_key(&mut self, raw_key: &[u8]) -> Result> { let item = self.snapshot.get(LOCK_CF, raw_key, 0); self.statistics.lock.get += 1; self.statistics.lock.flow_stats.read_keys += 1; self.statistics.lock.flow_stats.read_bytes += item.value_len(); self.statistics.lock.processed_keys += 1; self.statistics.total_size += raw_key.len() + item.value_len(); if item.value_len() == 0 { return Ok(None); } Ok(Some(txn_types::parse_lock(item.get_value())?)) } #[maybe_async::both] pub async fn get( &mut self, key: &Key, ts: TimeStamp, _gc_fence_limit: Option, ) -> Result> { let raw_key = key.to_raw()?; let item = self.snapshot.get(WRITE_CF, &raw_key, ts.into_inner()).await; self.statistics.write.get += 1; self.statistics.write.flow_stats.read_bytes += raw_key.len() + item.value_len(); self.statistics.write.flow_stats.read_keys += 1; self.statistics.total_size += raw_key.len() + item.value_len(); if item.value_len() > 0 { self.statistics.write.processed_keys += 1; self.statistics.processed_size += raw_key.len() + item.value_len(); return Ok(Some(item.get_value().to_vec())); } Ok(None) } #[maybe_async::both] pub async fn get_write( &mut self, key: &Key, ts: TimeStamp, _gc_fence_limit: Option, ) -> Result> { self.seek_write(key, ts) .await .map(|opt| opt.map(|(_, write)| write)) } #[maybe_async::both] pub async fn scan_write_for_key( &mut self, key: &Key, ts: TimeStamp, ) -> Result> { let raw_key = key.to_raw()?; let mut res = vec![]; let mut cursor = self .snapshot .new_iterator(WRITE_CF, false, true, None, self.fill_cache) .await; cursor.seek(&raw_key).await; while cursor.valid() { if cursor.key() != raw_key { break; } if !cursor.user_meta().is_empty() { let user_meta = UserMeta::from_slice(cursor.user_meta()); // only keep the write whose commit ts less or equal than required ts. if user_meta.commit_ts <= ts.into_inner() { let parsed = parse_write(&user_meta, cursor.val()); res.push(parsed); } } cursor.next().await; } Ok(res) } #[maybe_async::both] pub async fn seek_write( &mut self, key: &Key, ts: TimeStamp, ) -> Result> { let raw_key = key.to_raw()?; self.seek_write_with_raw_key(&raw_key, ts).await } #[maybe_async::both] pub async fn seek_write_with_raw_key( &mut self, raw_key: &[u8], ts: TimeStamp, ) -> Result> { let item = self.snapshot.get(WRITE_CF, raw_key, ts.into_inner()).await; self.statistics.write.seek += 1; self.statistics.write.flow_stats.read_keys += 1; self.statistics.write.flow_stats.read_bytes += raw_key.len() + item.value_len(); self.statistics.total_size += raw_key.len() + item.value_len(); if item.user_meta_len() > 0 { self.statistics.write.processed_keys += 1; self.statistics.processed_size += raw_key.len() + item.value_len(); let user_meta = UserMeta::from_slice(item.user_meta()); let (commit_ts, write) = parse_write(&user_meta, item.get_value()); return Ok(Some((commit_ts, write))); } Ok(None) } #[maybe_async::both] #[inline(always)] pub async fn get_old_value( &mut self, key: &Key, start_ts: TimeStamp, prev_write: Option, ) -> Result { if let Some(write) = prev_write { if write.write_type == WriteType::Delete { return Ok(OldValue::None); } // Locks and Rolbacks are stored in extra CF, will not be seeked by seek_write. assert_eq!(write.write_type, WriteType::Put); return Ok(OldValue::value(write.short_value.unwrap())); } let raw_key = key.to_raw()?; let item = self .snapshot .get(WRITE_CF, &raw_key, start_ts.into_inner()) .await; if item.value_len() > 0 { return Ok(OldValue::value(item.get_value().to_vec())); } Ok(OldValue::None) } fn create_lock_iter_on_range( &mut self, start: Option<&Key>, end: Option<&Key>, ) -> Result { let mut lock_iter = self.snapshot .new_iterator(LOCK_CF, false, false, None, self.fill_cache); let lower_bound: Bytes = if let Some(start) = start { Bytes::from(start.to_raw()?) } else { self.snapshot.clone_start_key() }; let upper_bound: Bytes = if let Some(end) = end { Bytes::from(end.to_raw()?) } else { self.snapshot.clone_end_key() }; if lock_iter.set_range(lower_bound, upper_bound) { self.statistics.lock.seek += 1; } self.statistics.lock.seek_tombstone += lock_iter.take_skipped_tombstones(); Ok(lock_iter) } fn lock_or_shared_locks_len(lock_or_shared_locks: &LockOrSharedLocks) -> usize { match lock_or_shared_locks { Either::Left(_) => 1, Either::Right(shared_locks) => shared_locks.len(), } } fn truncate_lock_or_shared_locks(lock_or_shared_locks: &mut LockOrSharedLocks, limit: usize) { if let Either::Right(shared_locks) = lock_or_shared_locks { shared_locks.truncate_to(limit); } } /// Scan locks that satisfies `filter(lock)` returns true, from the given /// start key `start`. At most `limit` locks will be returned. If /// `limit` is set to `0`, it means unlimited. /// /// Note: This only scans locks persisted in storage (LOCK_CF), not /// in-memory pessimistic locks. Use `SnapshotReader::scan_locks` for /// scanning both. /// /// The return type is `(locks, is_remain)`. `is_remain` indicates whether /// there MAY be remaining locks that can be scanned. #[maybe_async::both] pub async fn scan_locks_from_storage( &mut self, start: Option<&Key>, end: Option<&Key>, filter: F, limit: usize, ) -> Result<(Vec<(Key, LockOrSharedLocks)>, bool)> where F: Fn(&Lock) -> bool, { let mut locks = vec![]; let mut lock_count = 0usize; let mut lock_iter = self.create_lock_iter_on_range(start, end)?; // The initial seek already walks past every tombstone and filtered // version before the first live lock, so take the first checkpoint // before entering the loop. Entries consumed inside a single iterator // advance cannot be subdivided at this layer. let mut checkpointed_entries = 0usize; if self.take_lock_scan_checkpoint_due(&lock_iter, &mut checkpointed_entries) { check_lock_scan_reschedule().await; } let mut is_remain = false; while lock_iter.valid() { let key = Key::from_raw(lock_iter.key()); self.statistics.lock.flow_stats.read_keys += 1; self.statistics.lock.flow_stats.read_bytes += key.len(); if let Some(end) = end { if key >= *end { is_remain = false; break; } } let val = lock_iter.val(); self.statistics.lock.flow_stats.read_bytes += val.len(); let mut lock_or_shared_locks = txn_types::parse_lock(val)?; match &mut lock_or_shared_locks { Either::Left(lock) => { if filter(lock) { lock_count += 1; locks.push((key, lock_or_shared_locks)); } } Either::Right(shared_locks) => { shared_locks.filter_shared_locks(|lock| filter(lock))?; if !shared_locks.is_empty() { if limit > 0 { let remaining = limit.saturating_sub(lock_count); if remaining == 0 { is_remain = true; break; } if shared_locks.len() > remaining { shared_locks.truncate_to(remaining); is_remain = true; } } lock_count += shared_locks.len(); locks.push((key, lock_or_shared_locks)); } } } if limit > 0 && lock_count >= limit { is_remain = true; break; } lock_iter.next(); self.statistics.lock.next += 1; self.statistics.lock.next_tombstone += lock_iter.take_skipped_tombstones(); if self.take_lock_scan_checkpoint_due(&lock_iter, &mut checkpointed_entries) { check_lock_scan_reschedule().await; } } self.statistics.lock.processed_keys += lock_count; self.statistics.total_size += lock_iter.take_iter_stats().total_size; Ok((locks, is_remain)) } /// Scans locks from both storage and optionally in-memory pessimistic /// locks. #[maybe_async::both] pub async fn scan_locks( &mut self, in_memory_lock_ctx: Option, start: Option<&Key>, end: Option<&Key>, filter: F, limit: usize, ) -> Result<(Vec<(Key, LockOrSharedLocks)>, bool)> where F: Fn(&Lock) -> bool, { let (memory_locks, memory_has_remain) = match in_memory_lock_ctx { Some(ctx) => { let locks = ctx.txn_ext.pessimistic_locks.read(); locks .check_term_version_status(ctx.term, ctx.version) .map_err(|e| KvError::from(e))?; // Pass filter directly to scan_locks to avoid scanning unnecessary locks. // The filter is adapted to work with PessimisticLock by converting to Lock. locks.scan_locks(start, end, |_key, lock| filter(&lock.to_lock()), limit) } None => (Vec::new(), false), }; if memory_locks.is_empty() { return self .scan_locks_from_storage(start, end, filter, limit) .await; } let (storage_locks, storage_has_remain) = self .scan_locks_from_storage(start, end, &filter, limit) .await?; let capacity = if limit == 0 { memory_locks.len().saturating_add(storage_locks.len()) } else { limit.min(memory_locks.len().saturating_add(storage_locks.len())) }; let mut locks = Vec::with_capacity(capacity); let mut lock_count = 0usize; let mut memory_iter = memory_locks .into_iter() .map(|(key, lock)| (key, Either::Left(lock))) .peekable(); let mut storage_iter = storage_locks.into_iter().peekable(); let mut memory_pair = memory_iter.next(); let mut storage_pair = storage_iter.next(); let has_remain = loop { let mut next_pair = match (memory_pair.as_ref(), storage_pair.as_ref()) { (Some((memory_key, _)), Some((storage_key, _))) => { if storage_key <= memory_key { let next_key = storage_pair.take().unwrap(); storage_pair = storage_iter.next(); next_key } else { let next_key = memory_pair.take().unwrap(); memory_pair = memory_iter.next(); next_key } } (Some(_), None) => { let next_key = memory_pair.take().unwrap(); memory_pair = memory_iter.next(); next_key } (None, Some(_)) => { let next_key = storage_pair.take().unwrap(); storage_pair = storage_iter.next(); next_key } (None, None) => break memory_has_remain || storage_has_remain, }; // The same key could exist in both memory and storage when there is ongoing // leader transfer, split or merge on this region. In this case, duplicated // keys should be ignored. if locks.last().is_some_and(|(key, _)| key == &next_pair.0) { continue; } let next_count = Self::lock_or_shared_locks_len(&next_pair.1); if next_count == 0 { continue; } if limit > 0 { let remaining = limit.saturating_sub(lock_count); if remaining == 0 { break true; } if next_count > remaining { Self::truncate_lock_or_shared_locks(&mut next_pair.1, remaining); locks.push(next_pair); break true; } } lock_count += next_count; locks.push(next_pair); if limit > 0 && lock_count >= limit { break memory_pair.is_some() || storage_pair.is_some() || memory_has_remain || storage_has_remain; } }; Ok((locks, has_remain)) } /// Returns an arbitrary write that is newer than `ts`, or None if no such /// write exists. #[maybe_async::both] pub async fn get_newer_from_write_cf( &mut self, key: &Key, ts: TimeStamp, ) -> Result> { let raw_key = key.to_raw()?; self.get_newer_with_raw_key(&raw_key, ts).await } /// Returns an arbitrary write that is newer than `ts`, or None if no such /// write exists. #[maybe_async::both] pub async fn get_newer_with_raw_key( &mut self, raw_key: &[u8], ts: TimeStamp, ) -> Result> { let item = self .snapshot .get_newer(WRITE_CF, raw_key, ts.into_inner()) .await; if item.user_meta_len() > 0 { let user_meta = UserMeta::from_slice(item.user_meta()); let (commit_ts, write) = parse_write(&user_meta, item.get_value()); return Ok(Some((commit_ts, write))); } Ok(None) } /// Return the first committed key for which `start_ts` equals to `ts` // TODO: support async. pub fn seek_ts(&mut self, ts: TimeStamp) -> Result> { let mut it = self .snapshot .new_iterator(WRITE_CF, false, true, None, self.fill_cache); it.rewind(); while it.valid() { debug_assert!(!kvengine::table::is_deleted(it.meta())); // TODO: remove this check, iterator should not return deleted records. if kvengine::table::is_deleted(it.meta()) { it.next(); continue; } let user_meta = UserMeta::from_slice(it.user_meta()); if user_meta.start_ts == ts.into_inner() { return Ok(Some(Key::from_raw(it.key()))); } it.next() } Ok(None) } pub fn get_extras(&mut self, key: &Key) -> Result> { let raw_key = key.to_raw().unwrap(); let mut writes = vec![]; let mut extra_iter = self .snapshot .new_iterator(EXTRA_CF, false, false, None, true); extra_iter.seek(&raw_key); while extra_iter.valid() { if !extra_iter.key().starts_with(&raw_key) { break; } validate_extra_key_matches_raw_key(&raw_key, extra_iter.key())?; let um = UserMeta::from_slice(extra_iter.user_meta()); let (ts, write_type) = if um.is_rollback() { (um.start_ts, WriteType::Rollback) } else { (um.commit_ts, WriteType::Lock) }; writes.push((ts.into(), Write::new(write_type, um.start_ts.into(), None))); extra_iter.next(); } Ok(writes) } pub fn take_statistics(&mut self) -> tikv_kv::Statistics { std::mem::take(&mut self.statistics) } } pub fn parse_write(user_meta: &UserMeta, value: &[u8]) -> (TimeStamp, Write) { let commit_ts = user_meta.commit_ts; let write_type: WriteType; let short_value: Option; if value.is_empty() { write_type = WriteType::Delete; short_value = None; } else { write_type = WriteType::Put; short_value = Some(value.to_vec()) } ( TimeStamp::new(commit_ts), Write::new(write_type, TimeStamp::new(user_meta.start_ts), short_value), ) } #[cfg(test)] mod tests { use std::{ sync::{ Arc, atomic::{AtomicUsize, Ordering}, }, time::Duration, }; use cpu_time::ThreadTime; use kvengine::{ UserMeta, test_engine::{ DEF_BLOCK_SIZE, TABLE_KEY_PREFIX, install_lock_cf_test_table, new_test_engine_opt, }, }; use tikv_util::{ mpsc, yatp_pool::{ DefaultTicker, FairAdmitOutcome, FairCpuAdmissionMetadata, FairCpuBudgetAction, FairCpuPermit, FairQueueAdmissionController, FairQueueAdmissionResult, TaskMetadata, YatpPoolBuilder, set_task_meta, }, }; use txn_types::{Key, Lock, LockType, TimeStamp}; use yatp::{queue::Extras, task::future::TaskCell}; use super::{CloudReader, ExtraCfScanStrategyKind}; use crate::{read_pool::FairCpuTaskType, storage::config::ReadRescheduleOptions}; 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)) } } fn lock_bytes(key: &[u8], start_ts: u64) -> Vec { Lock::new( LockType::Put, key.to_vec(), TimeStamp::new(start_ts), 3000, Some(b"value".to_vec()), TimeStamp::zero(), 1, TimeStamp::zero(), ) .to_bytes() } fn read_reschedule_options(checkpoint_entries: usize) -> ReadRescheduleOptions { ReadRescheduleOptions { cloud_store_fair_cpu_checkpoint_entries: checkpoint_entries, ..Default::default() } } fn run_fair_scan( mut reader: CloudReader, start: Key, end: Key, filter: fn(&Lock) -> bool, limit: usize, ) -> ( super::Result<(Vec<(Key, txn_types::LockOrSharedLocks)>, bool)>, tikv_kv::Statistics, usize, ) { let controller = Arc::new(YieldingAdmissionController::default()); let mut builder = YatpPoolBuilder::new(DefaultTicker); builder .name_prefix("test-cloud-reader-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 fair-CPU checkpoint's one-millisecond sampling interval // due before the entry cadence reaches its first safe point. { let timer = ThreadTime::now(); while timer.elapsed() < Duration::from_millis(2) { std::hint::black_box(()); } } let result = reader .scan_locks_from_storage_async(Some(&start), Some(&end), filter, limit) .await; tx.send((result, reader.take_statistics())).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: FairCpuTaskType::KvScanLock as u8, deadline: None, }), }, ); pool.spawn(TaskCell::new(future, extras)); let (result, stats) = rx.recv_timeout(Duration::from_secs(10)).unwrap(); (result, stats, controller.admissions.load(Ordering::SeqCst)) } #[test] #[should_panic(expected = "legacy rollback flag true != pruned rollback flag false")] fn test_both_mode_compare_panics_on_rollback_kind_mismatch() { CloudReader::compare_both_mode_results( ExtraCfScanStrategyKind::unified, Some(UserMeta::new(10, 0)), Some(UserMeta::new(11, 12)), ); } #[test] fn test_scan_locks_async_yields_after_filtered_entries_without_changing_results() { const LOCK_COUNT: usize = 128; const CHECKPOINT_ENTRIES: usize = 16; const FIRST_KEY_INDEX: usize = 300; let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let entries = (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, 2001, &entries); let start = Key::from_raw(&kb.i_to_outer_key(FIRST_KEY_INDEX)); let end = Key::from_raw(&kb.i_to_outer_key(FIRST_KEY_INDEX + LOCK_COUNT)); let snapshot = engine.get_shard(1).unwrap().new_snap_access(); // Establish the synchronous behavior for a range whose locks are all // examined but rejected by the caller's max-ts-style filter. let mut sync_reader = CloudReader::new(snapshot.clone(), false) .with_read_reschedule_options(read_reschedule_options(CHECKPOINT_ENTRIES)); let (sync_locks, sync_has_remain) = sync_reader .scan_locks_from_storage(Some(&start), Some(&end), |_| false, 1) .unwrap(); let sync_stats = sync_reader.take_statistics(); assert!(sync_locks.is_empty()); assert!(!sync_has_remain); assert_eq!(sync_stats.lock.next, LOCK_COUNT); assert_eq!(sync_stats.lock.processed_keys, 0); let async_reader = CloudReader::new(snapshot, false) .with_read_reschedule_options(read_reschedule_options(CHECKPOINT_ENTRIES)); let (result, async_stats, admissions) = run_fair_scan(async_reader, start, end, |_| false, 1); let (async_locks, async_has_remain) = result.unwrap(); assert!(async_locks.is_empty()); assert_eq!(async_has_remain, sync_has_remain); assert_eq!(async_stats.lock.next, sync_stats.lock.next); assert_eq!( async_stats.lock.processed_keys, sync_stats.lock.processed_keys ); assert!( admissions > 1, "physical LOCK_CF entries must checkpoint even when the filter returns no locks" ); } #[test] fn test_scan_locks_async_checkpoints_tombstones_skipped_by_initial_seek() { const TOMBSTONE_COUNT: usize = 64; const CHECKPOINT_ENTRIES: usize = 16; const FIRST_KEY_INDEX: usize = 500; let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let mut entries = (0..TOMBSTONE_COUNT) .map(|offset| (FIRST_KEY_INDEX + offset, Vec::new())) .collect::>(); let live_key_index = FIRST_KEY_INDEX + TOMBSTONE_COUNT; let live_key = kb.i_to_outer_key(live_key_index); entries.push((live_key_index, lock_bytes(&live_key, 10))); install_lock_cf_test_table(&engine, 2002, &entries); let start = Key::from_raw(&kb.i_to_outer_key(FIRST_KEY_INDEX)); let end = Key::from_raw(&kb.i_to_outer_key(live_key_index + 1)); let snapshot = engine.get_shard(1).unwrap().new_snap_access(); let reader = CloudReader::new(snapshot, false) .with_read_reschedule_options(read_reschedule_options(CHECKPOINT_ENTRIES)); // Stop at the first live lock. With this limit, the loop does not call // `next` after processing it, so only the pre-loop checkpoint can // account for the tombstones skipped by the initial seek. let (result, stats, admissions) = run_fair_scan(reader, start, end, |_| true, 1); let (locks, has_remain) = result.unwrap(); assert_eq!(locks.len(), 1); assert_eq!(locks[0].0, Key::from_raw(&live_key)); assert!(has_remain); assert_eq!(stats.lock.seek_tombstone, TOMBSTONE_COUNT); assert_eq!(stats.lock.processed_keys, 1); assert!( admissions > 1, "the initial seek must checkpoint after skipping many LOCK_CF tombstones" ); } }