// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0. // Modifications Copyright 2022 PingCAP, Inc. use std::{ borrow::Cow, cell::RefCell, collections::HashMap, fmt::{self, Display}, future::Future, iter::FromIterator, marker::PhantomData, mem, ops::Deref, sync::{ Arc, Mutex, atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, }, time::Duration, }; use ::tracker::{ GLOBAL_TRACKERS, INVALID_TRACKER_TOKEN, RequestInfo, RequestType, TrackedFuture, TrackerToken, get_tls_tracker_token, set_tls_tracker_token, with_tls_tracker, }; use anyhow::anyhow; use api_version::{KvFormat, dispatch_api_version}; use async_stream::try_stream; use concurrency_manager::ConcurrencyManager; use futures::{ channel::{mpsc, oneshot}, prelude::*, }; use futures_util::future::try_join_all; use kvengine::{ KvEngineConfig, SnapAccess, context::{IaCtx, SnapCtx}, }; use kvproto::{ coprocessor as coppb, errorpb, kvrpcpb, kvrpcpb::{ScanDetailV2, TimeDetail, TimeDetailV2}, metapb, }; use online_config::ConfigManager; use overload_protector::{CopTaskStats, OverloadProtector}; use protobuf::{CodedInputStream, Message}; use raftstore::{RegionInfoAccessor, coprocessor::RegionInfoProvider}; use resource_control::{ CPU_THROTTLE_LOW_PRIORITY_DAG_CPU_TOKEN_WAIT_DURATION_HISTOGRAM, CPU_THROTTLE_LOW_PRIORITY_DAG_KEYSPACE_WAIT_DURATION_HISTOGRAM, CPU_THROTTLE_LOW_PRIORITY_DAG_REQUEST_ACTUAL_TO_ESTIMATED_RATIO_HISTOGRAM, CPU_THROTTLE_LOW_PRIORITY_DAG_REQUEST_CPU_TIME_HISTOGRAM, CPU_THROTTLE_LOW_PRIORITY_DAG_WAIT_DURATION_HISTOGRAM, CPU_THROTTLE_REQUEST_ACTUAL_TO_ESTIMATED_RATIO_HISTOGRAM, CPU_THROTTLE_REQUEST_CPU_TIME_HISTOGRAM, CPU_THROTTLE_REQUEST_DAG_MAIN_CPU_TIME_HISTOGRAM, CPU_TOKEN_RETRY_COUNT, CPU_TOKEN_WAIT_DURATION_HISTOGRAM, CpuTokenCheckFuture, CpuTokenError, CpuTokenHandle, KeyspaceReadLimiter, Metric, REQUEST_WAIT_HISTOGRAM_VEC, ReadLimiter, Resource, ResourceController, ResourceEvent, ResourcePublisher, Scope, Severity, ThrottleError, }; use resource_metering::{ FutureExt, ResourceTagFactory, StreamExt, record_logical_read_bytes, record_network_in_bytes, record_network_out_bytes, }; use security::SecurityManager; use tidb_query_common::{ error::StorageError, execute_stats::ExecSummary, storage::{FindRegionResult, RegionStorageAccessor, Result as StorageResult}, }; use tikv_alloc::trace::MemoryTraceGuard; use tikv_kv::{ExtraRegionOverride, SnapshotExt}; use tikv_util::{ codec::bytes::encode_bytes, config::VersionTrack, cpu_observe::CpuObserveFuture, deadline::set_deadline_exceeded_busy_error, memory::{MemoryQuota, OwnedAllocated}, metrics::{ServerIsBusyType, new_server_is_busy}, quota_limiter::QuotaLimiter, store::find_peer, sys::SysQuota, time::{Instant, duration_to_us}, }; use tipb::{AnalyzeReq, AnalyzeType, ChecksumRequest, ChecksumScanOn, DagRequest, ExecType}; use tokio::sync::Semaphore; use txn_types::Lock; use crate::{ coprocessor::{ Error, cache::CachedRequestHandler, config_manager::CopConfigManager, interceptors::*, keyspace_cpu::report_keyspace_cpu_time, metrics::*, remote_dispatcher::{ FallbackAdmission, FallbackAdmissionFuture, FallbackLocalAdmission, RemoteContext, RemoteRequest, RemoteThenLocalFallbackDagHandler, estimate_remote_dag_dispatch_extra_bytes, estimate_remote_db9_dispatch_extra_bytes, remote_cop_request_tag, try_remote_dag_handler, try_remote_db9_handler, unsupported_remote_db9_request_reason, }, runtime_budget::CpuTokenRuntimeBudget, scan_limiter::{ScanLimiterRef, ScanRateLimiter}, tracker::Tracker, *, }, read_pool::{ FairCpuTaskType, ReadPoolHandle, ReadPoolTaskOptions, ReadPoolTaskPollCpuRecorder, sync_read_pool_metrics_to_current_tracker, }, server::Config, storage::{ self, Engine, Snapshot, config::{ReadRescheduleConfig, ReadRescheduleOptions}, kv::{self, SnapContext, with_tls_engine}, mvcc::Error as MvccError, need_check_locks, need_check_locks_in_replica_read, txn::CloudStore, }, }; /// Requests that need less than this amount of execution time are considered /// light and can finish without acquiring a concurrency permit. const LIGHT_TASK_THRESHOLD: Duration = Duration::from_millis(5); // Minimum interval between warning logs about additional CPU token checks. const ADDITIONAL_CPU_TOKEN_WARN_INTERVAL_SECS: u64 = 5; // Last warning log timestamp used to rate-limit warning log emission. static LAST_ADDITIONAL_CPU_TOKEN_WARN_TIME_SECS: AtomicU64 = AtomicU64::new(0); #[derive(Clone)] pub struct CpuThrottleMetricsContext { pub(crate) keyspace_id: u32, pub(crate) region_id: u64, pub(crate) peer: Option, pub(crate) request_type: RequestType, pub(crate) is_remote: bool, pub(crate) is_low_priority_dag: bool, pub(crate) estimated_cpu_us: Option, pub(crate) resource_publisher: Option>, pub(crate) resource_debug: bool, } pub(crate) async fn run_with_cpu_token_check( future: F, semaphore: Option<&Semaphore>, cpu_token_handle: Option>, resource_debug: bool, keyspace_id: u32, region_id: u64, ) -> Result<(Duration, T)> where F: Future, { let runtime_token_config = cpu_token_handle .as_ref() .and_then(|h| h.get_runtime_config().map(|cfg| (h.clone(), cfg))); let (cpu_time, result) = if let Some((handle, (check_interval, threshold, per_allocation_us))) = runtime_token_config { if resource_debug { info!( "[CPU throttle] Token runtime management enabled for keyspace {}", keyspace_id; "check_interval_us" => check_interval.as_micros(), "threshold" => threshold, "per_allocation_us" => per_allocation_us, ); } let check_future = CpuTokenCheckFuture::new( future, handle.clone(), check_interval, threshold, per_allocation_us, ); let (cpu_time, check_result) = if let Some(semaphore) = semaphore { CpuObserveFuture::new(limit_concurrency( check_future, semaphore, LIGHT_TASK_THRESHOLD, )) .await } else { CpuObserveFuture::new(check_future).await }; match check_result { Ok(inner_result) => (cpu_time, inner_result), Err(CpuTokenError::Timeout) => { tikv_util::warn!( "[CPU throttle] Request timeout waiting for CPU tokens"; "keyspace_id" => keyspace_id, "region_id" => region_id, "allocated_us" => handle.allocated(), ); return Err(Error::CpuThrottled( "[CPU throttle] Request timeout waiting for CPU tokens".to_string(), )); } } } else if let Some(semaphore) = semaphore { CpuObserveFuture::new(limit_concurrency(future, semaphore, LIGHT_TASK_THRESHOLD)).await } else { CpuObserveFuture::new(future).await }; Ok((cpu_time, result)) } pub(crate) fn finalize_cpu_throttle_metrics( metrics: &CpuThrottleMetricsContext, cpu_token_handle: &Arc, cpu_time: Duration, record_actual_usage: bool, ) { let cpu_elapsed_us = cpu_time.as_micros() as u64; if record_actual_usage { cpu_token_handle.record_actual_usage(cpu_elapsed_us); } let keyspace_id_str = metrics.keyspace_id.to_string(); CPU_THROTTLE_REQUEST_DAG_MAIN_CPU_TIME_HISTOGRAM .with_label_values(&[&keyspace_id_str]) .observe(cpu_elapsed_us as f64 / 1_000_000.0); if record_actual_usage { CPU_THROTTLE_REQUEST_CPU_TIME_HISTOGRAM .with_label_values(&[&keyspace_id_str]) .observe(cpu_elapsed_us as f64 / 1_000_000.0); } if metrics.is_low_priority_dag { CPU_THROTTLE_LOW_PRIORITY_DAG_REQUEST_CPU_TIME_HISTOGRAM .observe(cpu_elapsed_us as f64 / 1_000_000.0); } if let Some(estimated_cpu_us) = metrics.estimated_cpu_us { if estimated_cpu_us > 0 { let ratio = cpu_elapsed_us as f64 / estimated_cpu_us as f64; CPU_THROTTLE_REQUEST_ACTUAL_TO_ESTIMATED_RATIO_HISTOGRAM .with_label_values(&[&keyspace_id_str]) .observe(ratio); if metrics.is_low_priority_dag { CPU_THROTTLE_LOW_PRIORITY_DAG_REQUEST_ACTUAL_TO_ESTIMATED_RATIO_HISTOGRAM .observe(ratio); } } } // Final check: allocate more if needed (non-blocking). if record_actual_usage && cpu_elapsed_us > cpu_token_handle.allocated() { let additional_us = cpu_elapsed_us - cpu_token_handle.allocated(); if !cpu_token_handle.allocate_more(additional_us) { let now_secs = Instant::now_coarse().second().max(0) as u64; if should_log_additional_cpu_token_warn(now_secs) { tikv_util::warn!( "[CPU throttle] Failed to allocate additional CPU tokens after execution"; "keyspace_id" => metrics.keyspace_id, "allocated_us" => cpu_token_handle.allocated(), "actual_us" => cpu_elapsed_us, "additional_needed_us" => additional_us, ); } } } // Publish keyspace CPU metric before dropping the handle. if !record_actual_usage { return; } let actual_cpu_us = cpu_token_handle.actual_used(); let cpu_time_secs = actual_cpu_us as f64 / 1_000_000.0; if let Some(resource_publisher) = &metrics.resource_publisher { if !metrics.is_remote && matches!(metrics.request_type, RequestType::CoprocessorDag) && actual_cpu_us > 0 { if metrics.resource_debug { info!( "[CPU throttle] Publishing keyspace CPU metric"; "keyspace_id" => metrics.keyspace_id, "region_id" => metrics.region_id, "peer" => ?metrics.peer, "request_type" => ?metrics.request_type, "is_remote" => metrics.is_remote, "actual_cpu_us" => actual_cpu_us, "cpu_time_secs" => cpu_time_secs, ); } resource_publisher.publish(ResourceEvent::CollectMetric(Metric { resource: Resource::Cpu { cpu_type: resource_control::CpuType::UnifiedReadPool, cpu_time: cpu_time_secs, }, scope: Scope::Keyspace { keyspace_id: metrics.keyspace_id, }, severity: Severity::Normal, })); } } } pub(crate) fn prepare_cpu_token_runtime_check( handler: &mut dyn RequestHandler, cpu_token_handle: Option>, resource_debug: bool, keyspace_id: u32, region_id: u64, ) -> Option> { let handle = cpu_token_handle?; if handle.get_runtime_config().is_some() && handler .set_copr_runtime_budget(Some(Arc::new(CpuTokenRuntimeBudget::new(handle.clone())))) { if resource_debug { info!( "[CPU throttle] DAG batch-level runtime management enabled"; "keyspace_id" => keyspace_id, "region_id" => region_id, ); } None } else { Some(handle) } } /// Coprocessor runs on the cop-worker is much cheaper than on the tikv-server. /// Modify the size to reduce RU consumption. const SIZE_DISCOUNT_PERCENT: usize = 50; fn is_db9_index_scan_kind(scan_kind: kvproto::db9_coprocessor::Db9ScanKind) -> bool { matches!( scan_kind, kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindIndex | kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindIndexRange | kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindIndexBoundedRange | kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindInList ) } fn db9_req_tag_from_scan_kind(scan_kind: kvproto::db9_coprocessor::Db9ScanKind) -> ReqTag { if is_db9_index_scan_kind(scan_kind) { ReqTag::db9_index } else { ReqTag::db9_select } } /// Fixed, low-cardinality top-level type of a remote coprocessor request. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RemoteCopRequestType { Dag, Db9, Analyze, Checksum, Unknown, } /// Returns the fixed, low-cardinality top-level type of a remote coprocessor /// request without decoding its nested payload. pub fn remote_cop_request_type(req: &coppb::Request) -> RemoteCopRequestType { match req.get_tp() { REQ_TYPE_DAG => RemoteCopRequestType::Dag, REQ_TYPE_DB9_DAG => RemoteCopRequestType::Db9, REQ_TYPE_ANALYZE => RemoteCopRequestType::Analyze, REQ_TYPE_CHECKSUM => RemoteCopRequestType::Checksum, _ => RemoteCopRequestType::Unknown, } } #[inline] fn can_handle_by_cached_request(snap: &S, req_ctx: &ReqContext) -> bool { req_ctx.cache_match_version.is_some() && req_ctx.cache_match_version == snap.ext().get_data_version() } #[inline] fn is_low_priority_dag_request( read_limiter: &ReadLimiter, request_type: RequestType, priority: kvrpcpb::CommandPri, ) -> bool { matches!(request_type, RequestType::CoprocessorDag) && matches!(priority, kvrpcpb::CommandPri::Low) && read_limiter.is_low_priority_dag_cpu_throttle_enabled() } thread_local! { static LOCAL_COPR_CPU_STATS: RefCell = RefCell::new(ThreadLocalKeyspaceCpuStats::new()); } #[inline] fn should_log_additional_cpu_token_warn(now_secs: u64) -> bool { let mut last = LAST_ADDITIONAL_CPU_TOKEN_WARN_TIME_SECS.load(Ordering::Relaxed); loop { if now_secs.saturating_sub(last) < ADDITIONAL_CPU_TOKEN_WARN_INTERVAL_SECS { return false; } match LAST_ADDITIONAL_CPU_TOKEN_WARN_TIME_SECS.compare_exchange_weak( last, now_secs, Ordering::Relaxed, Ordering::Relaxed, ) { Ok(_) => return true, Err(current) => last = current, } } } struct CoprReadPoolPollCpuRecorder { keyspace_id: u32, region_id: u64, resource_debug: bool, token_handle: Mutex>>, poll_cpu_us: AtomicU64, finished: AtomicBool, resource_publisher: Arc, } impl fmt::Debug for CoprReadPoolPollCpuRecorder { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("CoprReadPoolPollCpuRecorder") .field("keyspace_id", &self.keyspace_id) .field("region_id", &self.region_id) .field("poll_cpu_us", &self.poll_cpu_us.load(Ordering::Acquire)) .finish_non_exhaustive() } } impl CoprReadPoolPollCpuRecorder { fn new( keyspace_id: u32, region_id: u64, resource_debug: bool, resource_publisher: Arc, ) -> Self { Self { keyspace_id, region_id, resource_debug, token_handle: Mutex::new(None), poll_cpu_us: AtomicU64::new(0), finished: AtomicBool::new(false), resource_publisher, } } fn set_token_handle(&self, token_handle: Arc) { *self.token_handle.lock().unwrap() = Some(token_handle); } } impl ReadPoolTaskPollCpuRecorder for CoprReadPoolPollCpuRecorder { fn observe_poll_cpu_time(&self, cpu_time: Duration) { self.poll_cpu_us .fetch_add(duration_to_us(cpu_time), Ordering::AcqRel); } fn finish(&self) { if self.finished.swap(true, Ordering::AcqRel) { return; } let poll_cpu_us = self.poll_cpu_us.load(Ordering::Acquire); if poll_cpu_us == 0 { return; } let keyspace_id = self.keyspace_id.to_string(); CPU_THROTTLE_REQUEST_CPU_TIME_HISTOGRAM .with_label_values(&[&keyspace_id]) .observe(poll_cpu_us as f64 / 1_000_000.0); if let Some(token_handle) = self.token_handle.lock().unwrap().clone() { token_handle.record_actual_usage(poll_cpu_us); let actual_cpu_us = token_handle.actual_used(); if actual_cpu_us > token_handle.allocated() { let additional_us = actual_cpu_us - token_handle.allocated(); if !token_handle.allocate_more(additional_us) { let now_secs = Instant::now_coarse().second().max(0) as u64; if should_log_additional_cpu_token_warn(now_secs) { tikv_util::warn!( "[CPU throttle] Failed to allocate additional CPU tokens after read-pool task execution"; "keyspace_id" => self.keyspace_id, "allocated_us" => token_handle.allocated(), "actual_us" => actual_cpu_us, "poll_cpu_us" => poll_cpu_us, "additional_needed_us" => additional_us, ); } } } } let cpu_time_secs = poll_cpu_us as f64 / 1_000_000.0; if self.resource_debug { info!( "[CPU throttle] Publishing read-pool task poll CPU metric"; "keyspace_id" => self.keyspace_id, "region_id" => self.region_id, "actual_cpu_us" => poll_cpu_us, "cpu_time_secs" => cpu_time_secs, ); } self.resource_publisher .publish(ResourceEvent::CollectMetric(Metric { resource: Resource::Cpu { cpu_type: resource_control::CpuType::UnifiedReadPool, cpu_time: cpu_time_secs, }, scope: Scope::Keyspace { keyspace_id: self.keyspace_id, }, severity: Severity::Normal, })); } } /// A pool to build and run Coprocessor request handlers. #[derive(Clone)] pub struct Endpoint { /// The thread pool to run Coprocessor requests. read_pool: ReadPoolHandle, /// The pre-fairness global concurrency limiter for long coprocessor tasks. /// It remains active on the non-fair YATP path to preserve existing /// behavior, while the fair path uses its keyspace-aware task limiter. semaphore: Option>, /// The memory quota for coprocessor requests. memory_quota: Arc, concurrency_manager: ConcurrencyManager, region_info_accessor: Option, resource_tag_factory: ResourceTagFactory, /// The recursion limit when parsing Coprocessor Protobuf requests. recursion_limit: u32, batch_row_limit: usize, stream_batch_row_limit: usize, stream_channel_size: usize, /// The soft time limit of handling Coprocessor requests. max_handle_duration: Duration, slow_log_threshold: Duration, _phantom: PhantomData, max_resp_size: Arc, remote_ctx: Option, remote_pool: Option>, /// Pre-fairness request admission controls. Both are `None` when /// keyspace fairness is enabled, making the bypass independent of online /// legacy configuration. quota_limiter: Option>, overload_protector: Option, security_mgr: Arc, pub resource_publisher: Option>, pub resource_controller: Option, pub read_limiter: Option, status_addr: String, read_reschedule_config: ReadRescheduleConfig, } /// The result of parsing a Coprocessor request. pub struct ParseCopRequestResult { req_tag: ReqTag, req_ctx: ReqContext, remote_dispatch_precharge_bytes: usize, handler_builder: RequestHandlerBuilder, } type SharedOwnedAllocated = Arc>; impl ParseCopRequestResult { #[cfg(test)] pub fn default_for_test(handler_builder: RequestHandlerBuilder) -> Self { Self { req_tag: ReqTag::test, req_ctx: ReqContext::default_for_test(), remote_dispatch_precharge_bytes: 0, handler_builder, } } } impl tikv_util::AssertSend for Endpoint {} impl Endpoint { async fn allocate_dag_cpu_tokens( req_ctx: &ReqContext, read_limiter: &ReadLimiter, keyspace_read_limiter: Option<&KeyspaceReadLimiter>, resource_dry_run: bool, resource_debug: bool, priority: kvrpcpb::CommandPri, ) -> Result<(Option>, Option)> { // This helper is for local DAG admission only. The call site is responsible // for gating on `!is_remote` and `request_type == CoprocessorDag`. let is_low_priority_dag = is_low_priority_dag_request(read_limiter, RequestType::CoprocessorDag, priority); let mut low_priority_keyspace_wait_duration = Duration::ZERO; if is_low_priority_dag { if let Some(keyspace_read_limiter) = keyspace_read_limiter { let wait_time = if resource_dry_run { keyspace_read_limiter.take_wait_time() } else { keyspace_read_limiter .wait_with_deadline(req_ctx.deadline.to_std_instant()) .await? }; REQUEST_WAIT_HISTOGRAM_VEC .with_label_values(&["read"]) .observe(wait_time.as_secs_f64()); low_priority_keyspace_wait_duration = wait_time; CPU_THROTTLE_LOW_PRIORITY_DAG_KEYSPACE_WAIT_DURATION_HISTOGRAM .observe(wait_time.as_secs_f64()); if resource_debug && !wait_time.is_zero() { info!( "resource control keyspace {}, region {}, task {}, mock {} sleep {:?} for DAG request", req_ctx.context.get_keyspace_id(), req_ctx.context.region_id, req_ctx.context.task_id, resource_dry_run, wait_time ); } } } let Some(cpu_throttle_manager) = read_limiter.get_cpu_throttle_manager() else { return Ok((None, None)); }; let keyspace_id = req_ctx.context.get_keyspace_id(); read_limiter.register_keyspace_for_service_limit(keyspace_id); let estimated_cpu_us = if is_low_priority_dag { cpu_throttle_manager.get_estimated_cpu_per_low_priority_dag_request_us(keyspace_id) } else { cpu_throttle_manager.get_estimated_cpu_per_request_us(keyspace_id) }; let request_deadline = req_ctx.deadline.to_std_instant(); let keyspace_id_str = keyspace_id.to_string(); let allocation_start = std::time::Instant::now(); let mut retry_count = 0u64; loop { let allocation_result = if is_low_priority_dag { cpu_throttle_manager.try_allocate_low_priority_dag( keyspace_id, estimated_cpu_us, request_deadline, ) } else { cpu_throttle_manager.try_allocate(keyspace_id, estimated_cpu_us, request_deadline) }; match allocation_result { Ok(handle) => { let handle = Arc::new(handle); let wait_duration = allocation_start.elapsed(); CPU_TOKEN_WAIT_DURATION_HISTOGRAM .with_label_values(&[&keyspace_id_str, "success"]) .observe(wait_duration.as_secs_f64()); if is_low_priority_dag { let total_wait_duration = low_priority_keyspace_wait_duration + wait_duration; CPU_THROTTLE_LOW_PRIORITY_DAG_CPU_TOKEN_WAIT_DURATION_HISTOGRAM .with_label_values(&["success"]) .observe(wait_duration.as_secs_f64()); CPU_THROTTLE_LOW_PRIORITY_DAG_WAIT_DURATION_HISTOGRAM .with_label_values(&["success"]) .observe(total_wait_duration.as_secs_f64()); } if retry_count > 0 { CPU_TOKEN_RETRY_COUNT .with_label_values(&[&keyspace_id_str]) .inc_by(retry_count); } if resource_debug { info!( "[CPU throttle] Token allocated for keyspace {}, estimated_cpu_us: {}, wait_duration: {:?}, retries: {}", keyspace_id, estimated_cpu_us, wait_duration, retry_count ); } return Ok((Some(handle), Some(estimated_cpu_us))); } Err(err) => { if resource_dry_run { if resource_debug { info!( "[CPU throttle] (dry_run): {:?} for keyspace {}", err, keyspace_id ); } // Match the old behavior: in dry-run mode we don't block, and we // don't treat the (expected) allocation failure as an admission event. return Ok((None, None)); } let now = std::time::Instant::now(); if now >= request_deadline { let wait_duration = allocation_start.elapsed(); let timeout_duration = request_deadline.saturating_duration_since(allocation_start); CPU_TOKEN_WAIT_DURATION_HISTOGRAM .with_label_values(&[&keyspace_id_str, "timeout"]) .observe(wait_duration.as_secs_f64()); if is_low_priority_dag { let total_wait_duration = low_priority_keyspace_wait_duration + wait_duration; CPU_THROTTLE_LOW_PRIORITY_DAG_CPU_TOKEN_WAIT_DURATION_HISTOGRAM .with_label_values(&["timeout"]) .observe(wait_duration.as_secs_f64()); CPU_THROTTLE_LOW_PRIORITY_DAG_WAIT_DURATION_HISTOGRAM .with_label_values(&["timeout"]) .observe(total_wait_duration.as_secs_f64()); } if retry_count > 0 { CPU_TOKEN_RETRY_COUNT .with_label_values(&[&keyspace_id_str]) .inc_by(retry_count); } if resource_debug { info!( "[CPU throttle] Request deadline exceeded"; "keyspace_id" => keyspace_id, "region_id" => req_ctx.context.region_id, "timeout_duration" => ?timeout_duration, "wait_duration" => ?wait_duration, "retry_count" => retry_count, "estimated_cpu_us" => estimated_cpu_us, "error" => ?err, ); } let err_msg = match err { ThrottleError::GlobalCpuExhausted => { "[CPU throttle] Global CPU quota exhausted (timeout waiting for tokens)".to_string() } ThrottleError::LowPriorityDagGlobalCpuExhausted => { "[CPU throttle] Low priority DAG global CPU quota exhausted (timeout waiting for tokens)".to_string() } ThrottleError::KeyspaceCpuExhausted => { format!( "[CPU throttle] Keyspace {} CPU quota exhausted (timeout waiting for tokens)", keyspace_id ) } ThrottleError::KeyspaceFairShareExceeded { .. } => { format!( "[CPU throttle] Keyspace {} exceeded fair share allocation (timeout waiting for tokens)", keyspace_id ) } ThrottleError::KeyspaceCpuExhaustedNoBurst { .. } => { format!( "[CPU throttle] Keyspace {} CPU quota exhausted and burst not allowed (timeout waiting for tokens)", keyspace_id ) } ThrottleError::RequestTimeout => { format!( "[CPU throttle] Keyspace {} request timeout during CPU token allocation", keyspace_id ) } }; return Err(Error::CpuThrottled(err_msg)); } let retry_at = cpu_throttle_manager .allocation_retry_at(keyspace_id, &err) .min(request_deadline); let wait_duration = retry_at.saturating_duration_since(now); if resource_debug { info!( "[CPU throttle] Token allocation failed for keyspace {}, waiting {:?} before retry: {:?}", keyspace_id, wait_duration, err ); } retry_count += 1; tikv_util::retry::sleep_async(wait_duration).await; } } } } pub fn new( cfg: &Config, read_pool: ReadPoolHandle, concurrency_manager: ConcurrencyManager, region_info_accessor: Option, resource_tag_factory: ResourceTagFactory, security_mgr: Arc, ) -> Self { let semaphore = read_pool .uses_coprocessor_concurrency_limiter() .then(|| Arc::new(Semaphore::new(cfg.end_point_max_concurrency))); let memory_quota = Arc::new(MemoryQuota::new(cfg.end_point_memory_quota.0 as _)); register_coprocessor_memory_quota_metrics(memory_quota.clone()); let status_addr = if cfg.advertise_status_addr.is_empty() { cfg.status_addr.clone() } else { cfg.advertise_status_addr.clone() }; Self { read_pool, semaphore, memory_quota, concurrency_manager, region_info_accessor, resource_tag_factory, recursion_limit: cfg.end_point_recursion_limit, batch_row_limit: cfg.end_point_batch_row_limit, stream_batch_row_limit: cfg.end_point_stream_batch_row_limit, stream_channel_size: cfg.end_point_stream_channel_size, max_handle_duration: cfg.end_point_request_max_handle_duration.0, slow_log_threshold: cfg.end_point_slow_log_threshold.0, _phantom: Default::default(), remote_ctx: None, remote_pool: None, quota_limiter: None, overload_protector: None, security_mgr, resource_publisher: None, resource_controller: None, read_limiter: None, status_addr, max_resp_size: Arc::new(AtomicU64::new(cfg.cop_max_resp_size.0)), read_reschedule_config: ReadRescheduleConfig::default(), } } /// Installs the startup-gated pre-fairness admission controls. pub fn set_legacy_admission_controls( &mut self, quota_limiter: Option>, overload_protector: Option, ) { self.quota_limiter = quota_limiter; self.overload_protector = overload_protector; } pub fn config_manager(&self) -> Box { Box::new(CopConfigManager::new( self.max_resp_size.clone(), self.memory_quota.clone(), )) } /// Builds the remote coprocessor context from the provided configuration. /// /// The context captures the current resource controller's total CPU /// monitor. Installing a resource controller after this call does not /// update the context. Use [`Self::configure_remote_coprocessor`] when /// configuring both. pub fn set_remote_url(&mut self, kvengine_config: Arc>) { let worker_threads = (SysQuota::cpu_cores_quota() as usize / 4).max(2); let pool = Arc::new( tokio::runtime::Builder::new_multi_thread() .worker_threads(worker_threads) .enable_all() .thread_name_fn(|| { static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0); let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst); format!("remote_coprocessor_{}", id) }) .build() .unwrap(), ); self.remote_ctx = Some(RemoteContext::new( kvengine_config, self.security_mgr.clone(), pool.handle().clone(), self.resource_controller .as_ref() .map(|controller| controller.total_cpu_monitor()), self.status_addr.clone(), )); self.remote_pool = Some(pool); } /// Configures the resource controller before building the remote /// coprocessor context. /// /// `RemoteContext` captures the controller's total CPU monitor when it is /// built, so the controller must be installed first for load-aware /// dispatch to receive CPU signals. pub fn configure_remote_coprocessor( &mut self, kvengine_config: Arc>, controller: ResourceController, ) { self.set_resource_controller(controller); self.set_remote_url(kvengine_config); } /// Shares storage's online read-reschedule settings with local /// coprocessor requests. pub fn set_read_reschedule_config(&mut self, config: ReadRescheduleConfig) { self.read_reschedule_config = config; } /// Installs the resource controller used by the coprocessor endpoint. /// /// This does not update an existing remote coprocessor context. Use /// [`Self::configure_remote_coprocessor`] when configuring both so the /// context captures the controller's total CPU monitor. pub fn set_resource_controller(&mut self, mut controller: ResourceController) { self.resource_publisher = Some(controller.publisher()); self.read_limiter = controller.get_read_limiter(); info!( "Coprocessor endpoint set resource controller, read_limiter: {}, debug: {}", self.read_limiter.is_some(), controller.get_debug() ); self.resource_controller = Some(controller); } fn check_memory_locks( concurrency_manager: &ConcurrencyManager, req_ctx: &ReqContext, ) -> Result<()> { check_memory_locks_for_ranges(concurrency_manager, req_ctx, &req_ctx.ranges) } fn parse_request_and_check_memory_locks( &self, req: coppb::Request, peer: Option, is_streaming: bool, max_resp_size: u64, ) -> Result> { let api_version = req.get_context().get_api_version(); dispatch_api_version!( api_version, { self.parse_request_and_check_memory_locks_impl::( req, peer, is_streaming, max_resp_size, ) }, Err(box_err!( "API V3 is not supported on the cloud-engine branch" )) ) } /// Parse the raw `Request` to create `RequestHandlerBuilder` and /// `ReqContext`. Returns `Err` if fails. /// /// It also checks if there are locks in memory blocking this read request. fn parse_request_and_check_memory_locks_impl( &self, mut req: coppb::Request, peer: Option, is_streaming: bool, max_resp_size: u64, ) -> Result> { fail_point!("coprocessor_parse_request", |_| Err(box_err!( "unsupported tp (failpoint)" ))); // Keep one request on a stable cadence if an online update arrives // while it is running. let read_reschedule_options = self.read_reschedule_config.snapshot(); let cop_req = if self .remote_ctx .as_ref() .and_then(RemoteContext::active_worker_addr) .is_some() && (req.get_tp() == REQ_TYPE_ANALYZE || req.get_tp() == REQ_TYPE_CHECKSUM) { req.write_to_bytes().unwrap() } else { Vec::default() }; let (context, data, ranges, mut start_ts) = ( req.take_context(), req.take_data(), req.take_ranges().to_vec(), req.get_start_ts(), ); let cache_match_version = if req.get_is_cache_enabled() { Some(req.get_cache_if_match_version()) } else { None }; let mut input = CodedInputStream::from_bytes(&data); input.set_recursion_limit(self.recursion_limit); let req_ctx: ReqContext; let req_tag: ReqTag; let builder: RequestHandlerBuilder; let mut remote_dispatch_precharge_bytes = 0; match req.get_tp() { REQ_TYPE_DAG => { let mut dag = DagRequest::default(); box_try!(dag.merge_from(&mut input)); let mut table_scan = false; let mut is_desc_scan = false; if let Some(scan) = dag.get_executors().iter().next() { table_scan = scan.get_tp() == ExecType::TypeTableScan; if table_scan { is_desc_scan = scan.get_tbl_scan().get_desc(); } else { is_desc_scan = scan.get_idx_scan().get_desc(); } } if start_ts == 0 { start_ts = dag.get_start_ts_fallback(); } req_tag = if table_scan { ReqTag::select } else { ReqTag::index }; let lazy_remote_pattern = self .remote_ctx .as_ref() .filter(|ctx| { ctx.cop_worker_provider .get(context.get_keyspace_id(), 0) .is_some() }) .and_then(|ctx| ctx.extract(context.get_keyspace_id(), &ranges, &dag)); req_ctx = ReqContext::new( context, ranges, self.max_handle_duration, peer, Some(is_desc_scan), start_ts.into(), cache_match_version, lazy_remote_pattern, ); let scan_limiter = self.build_scan_limiter(&req_ctx); with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorDag; tracker.req_info.start_ts = start_ts; }); Endpoint::::check_memory_locks(&self.concurrency_manager, &req_ctx)?; let batch_row_limit = self.get_batch_row_limit(is_streaming); let remote_ctx = self.remote_ctx.clone(); let paging_size = match req.get_paging_size() { 0 => None, i => Some(i), }; let paging_size_bytes = match req.get_paging_size_bytes() { 0 => None, i => Some(i), }; remote_dispatch_precharge_bytes = estimate_remote_dag_dispatch_extra_bytes( &dag, &req_ctx, remote_ctx.as_ref(), paging_size, ); let region_info_accessor = self.region_info_accessor.clone(); let concurrency_manager = self.concurrency_manager.clone(); let quota_limiter = self.quota_limiter.clone(); let fallback_read_limiter = self.read_limiter.clone(); let fallback_resource_publisher = self.resource_publisher.clone(); let fallback_fair_global_cpu_gate_active = self.can_use_fair_global_cpu_gate_for_dag(); let (fallback_resource_enabled, fallback_resource_dry_run, fallback_resource_debug) = self.resource_controller .as_ref() .map(|controller| { ( controller.get_enabled(), controller.get_dry_run(), controller.get_debug(), ) }) .unwrap_or_default(); builder = Box::new(move |snap, req_ctx| { let remote_handler = try_remote_dag_handler( snap.get_kvengine_snap(), &dag, req_ctx, remote_ctx, paging_size, ); let remote_tag = snap .get_kvengine_snap() .map(|snap| remote_cop_request_tag(snap, req_ctx)) .unwrap_or_else(|| "remote_dag".to_owned()); let local_req_ctx = req_ctx.clone(); let local_region_info_accessor = region_info_accessor.clone(); let local_concurrency_manager = concurrency_manager.clone(); let local_scan_limiter = scan_limiter.clone(); let local_quota_limiter = quota_limiter.clone(); let local_max_resp_size = max_resp_size; let local_paging_size = paging_size; let local_paging_size_bytes = paging_size_bytes; let local_is_cache_enabled = req.get_is_cache_enabled(); let local_batch_row_limit = batch_row_limit; let local_start_ts = start_ts; let local_builder = move || { let data_version = snap.ext().get_data_version(); let store = CloudStore::new( snap, local_start_ts, local_req_ctx.bypass_locks.clone(), !local_req_ctx.context.get_not_fill_cache(), // Only cache-enabled DAG requests need CloudStore to // track hidden newer MVCC data for can_be_cached. local_is_cache_enabled, ) .with_read_reschedule_options(read_reschedule_options); dag::DagHandlerBuilder::new( dag, local_req_ctx.ranges.clone(), store, ExtraSnapStoreAccessor::::new( local_req_ctx.clone(), local_region_info_accessor, local_concurrency_manager, read_reschedule_options, ), local_req_ctx.deadline, local_batch_row_limit, local_is_cache_enabled, local_paging_size, local_paging_size_bytes, local_max_resp_size, read_reschedule_options.range_scan(), ) .quota_limiter(local_quota_limiter) .scan_limiter(local_scan_limiter) .data_version(data_version) .build::() }; if let Some(handler) = remote_handler { with_tls_tracker(|tracker| { // Remote handlers must mark the request before returning so the // endpoint keeps the remote dispatch precharge instead of refunding // it as a local request. tracker.req_info.is_remote = true; }); let fallback_admission = Self::build_dag_fallback_admission( req_ctx, fallback_read_limiter.clone(), fallback_resource_publisher.clone(), fallback_resource_enabled, fallback_resource_dry_run, fallback_resource_debug, fallback_fair_global_cpu_gate_active, ); return Ok(RemoteThenLocalFallbackDagHandler::new( handler, Box::new(local_builder), fallback_admission, remote_tag, ) .into_boxed()); } local_builder() }); } REQ_TYPE_DB9_DAG => { let mut db9 = kvproto::db9_coprocessor::Db9DagRequest::default(); if let Err(err) = db9.merge_from(&mut input) { let err = Error::Other(err.to_string()); return Err(err); } req_tag = db9_req_tag_from_scan_kind(db9.get_scan().get_kind()); req_ctx = { ReqContextInner::new( context, ranges, self.max_handle_duration, peer, Some(db9.get_scan().get_desc()), start_ts.into(), // DB9 does not populate the normal DAG cache metadata, so // it must not use the generic request-cache fast path. None, None, ) .into() }; with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorDag; tracker.req_info.start_ts = start_ts; }); Endpoint::::check_memory_locks(&self.concurrency_manager, &req_ctx)?; let remote_ctx = self.remote_ctx.clone(); if !is_streaming { remote_dispatch_precharge_bytes = estimate_remote_db9_dispatch_extra_bytes( &db9, &req_ctx, remote_ctx.as_ref(), ); } let concurrency_manager = self.concurrency_manager.clone(); builder = Box::new(move |snap, req_ctx| { if !is_streaming { if let Some(handler) = try_remote_db9_handler( snap.get_kvengine_snap(), &db9, req_ctx, remote_ctx, ) { with_tls_tracker(|tracker| { // Remote handlers must mark the request before returning so the // endpoint keeps the remote dispatch precharge instead of refunding // it as a local request. tracker.req_info.is_remote = true; }); return Ok(handler); } } let extra_storage_accessor = ExtraSnapStoreAccessor::::new( req_ctx.clone(), None, concurrency_manager.clone(), read_reschedule_options, ); db9::Db9DagHandler::new( db9, snap, req_ctx, max_resp_size, extra_storage_accessor, read_reschedule_options, ) .map(|h| h.into_boxed()) }); } REQ_TYPE_ANALYZE => { let mut analyze = AnalyzeReq::default(); box_try!(analyze.merge_from(&mut input)); if start_ts == 0 { start_ts = analyze.get_start_ts_fallback(); } req_tag = match analyze.get_tp() { AnalyzeType::TypeIndex | AnalyzeType::TypeCommonHandle => ReqTag::analyze_index, AnalyzeType::TypeColumn | AnalyzeType::TypeMixed => ReqTag::analyze_table, AnalyzeType::TypeFullSampling => ReqTag::analyze_full_sampling, AnalyzeType::TypeSampleIndex => unimplemented!(), }; req_ctx = ReqContext::new( context, ranges, self.max_handle_duration, peer, None, start_ts.into(), cache_match_version, None, ); with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorAnalyze; tracker.req_info.start_ts = start_ts; }); Endpoint::::check_memory_locks(&self.concurrency_manager, &req_ctx)?; let remote_req = RemoteRequest { cop_req, ..Default::default() }; let remote_ctx = self.remote_ctx.clone(); let quota_limiter = self.quota_limiter.clone(); builder = Box::new(move |snap, req_ctx| { statistics::analyze::AnalyzeContext::<_, F>::new_with_quota( analyze, req_ctx.ranges.clone(), start_ts, snap, req_ctx, quota_limiter, remote_ctx, remote_req, read_reschedule_options, ) .map(|h| h.into_boxed()) }); } REQ_TYPE_CHECKSUM => { let mut checksum = ChecksumRequest::default(); box_try!(checksum.merge_from(&mut input)); let table_scan = checksum.get_scan_on() == ChecksumScanOn::Table; if start_ts == 0 { start_ts = checksum.get_start_ts_fallback(); } req_tag = if table_scan { ReqTag::checksum_table } else { ReqTag::checksum_index }; req_ctx = ReqContext::new( context, ranges, self.max_handle_duration, peer, None, start_ts.into(), cache_match_version, None, ); with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorChecksum; tracker.req_info.start_ts = start_ts; }); Endpoint::::check_memory_locks(&self.concurrency_manager, &req_ctx)?; let remote_req = RemoteRequest { cop_req, ..Default::default() }; let remote_ctx = self.remote_ctx.clone(); builder = Box::new(move |snap, req_ctx| { checksum::ChecksumContext::new( checksum, req_ctx.ranges.clone(), start_ts, snap, req_ctx, remote_ctx, remote_req, read_reschedule_options, ) .map(|h| h.into_boxed()) }); } tp => return Err(box_err!("unsupported tp {}", tp)), }; Ok(ParseCopRequestResult { req_tag, req_ctx, remote_dispatch_precharge_bytes, handler_builder: builder, }) } /// Get the batch row limit configuration. #[inline] fn get_batch_row_limit(&self, is_streaming: bool) -> usize { if is_streaming { self.stream_batch_row_limit } else { self.batch_row_limit } } fn build_scan_limiter(&self, req_ctx: &ReqContext) -> Option { self.resource_controller.as_ref().and_then(|controller| { if !controller.get_enabled() { return None; } let keyspace_id = req_ctx.context.get_keyspace_id(); let region_id = req_ctx.context.region_id; let task_id = req_ctx.context.task_id; let dry_run = controller.get_dry_run(); let debug = controller.get_debug(); self.read_limiter .as_ref() .and_then(|read_limiter| read_limiter.get_limiter(keyspace_id)) .map(|limiter| { Arc::new(ScanRateLimiter::new( limiter, req_ctx.deadline.to_std_instant(), dry_run, debug, keyspace_id, region_id, task_id, )) as ScanLimiterRef }) }) } fn add_ranges_to_snap_context(ctx: &ReqContext, snap_ctx: &mut SnapContext<'_>) { // Need to pass start_ts and ranges to check memory locks for replica read for r in &ctx.ranges { let start_key = txn_types::Key::from_raw(r.get_start()); let end_key = txn_types::Key::from_raw(r.get_end()); let mut key_range = kvrpcpb::KeyRange::default(); key_range.set_start_key(start_key.into_encoded()); key_range.set_end_key(end_key.into_encoded()); snap_ctx.key_ranges.push(key_range); } } #[inline] fn async_snapshot( engine: &mut E, req_ctx: &ReqContext, ) -> impl std::future::Future> { let mut snap_ctx = SnapContext { pb_ctx: &req_ctx.context, start_ts: Some(req_ctx.txn_start_ts), ..Default::default() }; if need_check_locks_in_replica_read(&req_ctx.context) { Self::add_ranges_to_snap_context(req_ctx, &mut snap_ctx); } kv::snapshot(engine, snap_ctx).map_err(Error::from) } fn build_dag_fallback_admission( req_ctx: &ReqContext, read_limiter: Option, resource_publisher: Option>, resource_enabled: bool, resource_dry_run: bool, resource_debug: bool, fair_global_cpu_gate_active: bool, ) -> Option { if !resource_enabled || !read_limiter .as_ref() .is_some_and(ReadLimiter::is_legacy_read_limiter_allowed) { return None; } let keyspace_read_limiter = read_limiter .as_ref() .and_then(|limiter| limiter.get_limiter(req_ctx.context.get_keyspace_id())); let req_ctx = req_ctx.clone(); Some(Box::new(move || { Box::pin(async move { let (handle, estimated_cpu_us_opt) = if fair_global_cpu_gate_active { (None, None) } else if let Some(read_limiter) = read_limiter.as_ref() { Self::allocate_dag_cpu_tokens( &req_ctx, read_limiter, keyspace_read_limiter.as_ref(), resource_dry_run, resource_debug, req_ctx.context.get_priority(), ) .await? } else { (None, None) }; Ok(handle.map(|cpu_token_handle| FallbackLocalAdmission { cpu_token_handle, keyspace_id: req_ctx.context.get_keyspace_id(), region_id: req_ctx.context.region_id, peer: req_ctx.context.peer.clone().into_option(), estimated_cpu_us: estimated_cpu_us_opt, is_low_priority_dag: read_limiter.as_ref().is_some_and(|limiter| { is_low_priority_dag_request( limiter, RequestType::CoprocessorDag, req_ctx.context.get_priority(), ) }), resource_publisher, resource_debug, })) }) as FallbackAdmissionFuture })) } fn can_use_fair_global_cpu_gate_for_dag(&self) -> bool { self.read_pool.uses_fair_global_cpu_gate() && self .resource_controller .as_ref() .is_some_and(|controller| controller.get_enabled()) && self .read_limiter .as_ref() .is_some_and(|limiter| limiter.is_cpu_throttle_enabled()) && self.resource_publisher.is_some() } #[inline] pub fn async_snapshot_for_extra( engine: &mut E, req_ctx: &ReqContext, extra: ExtraRegionOverride, ) -> impl std::future::Future> { let snap_ctx = SnapContext { pb_ctx: &req_ctx.context, start_ts: Some(req_ctx.txn_start_ts), extra_region_override: Some(extra), ..Default::default() }; kv::snapshot(engine, snap_ctx).map_err(Error::from) } /// The real implementation of handling a unary request. /// /// It first retrieves a snapshot, then builds the `RequestHandler` over the /// snapshot and the given `handler_builder`. Finally, it calls the unary /// request interface of the `RequestHandler` to process the request and /// produce a result. async fn handle_unary_request_impl( semaphore: Option>, mut tracker: Box>, handler_builder: RequestHandlerBuilder, remote_dispatch_precharge_bytes: usize, request_memory_quota_guard: Option, remote_ctx: Option, tracker_token: TrackerToken, resource_publisher: Option>, poll_cpu_recorder: Option>, fair_global_cpu_gate_active: bool, read_limiter: Option, keyspace_read_limiter: Option, resource_dry_run: bool, resource_debug: bool, ) -> Result> { with_tls_tracker(|tracker1| { record_network_in_bytes(tracker1.metrics.grpc_req_size); }); // When this function is being executed, it may be queued for a long time, so // that deadline may exceed. tracker.on_scheduled(); tracker.req_ctx.deadline.check()?; // Safety: spawning this function using a `FuturePool` ensures that a TLS engine // exists. let snapshot = unsafe { with_tls_engine(|engine| Self::async_snapshot(engine, &tracker.req_ctx)) } .await?; let latest_buckets = snapshot.ext().get_buckets(); // Check if the buckets version is latest. if let Some(ref buckets) = latest_buckets { if buckets.version > tracker.req_ctx.context.buckets_version { // check if the request has large range, we don't need to return error if the // range is small. let lower_bound_key = encode_bytes(&tracker.req_ctx.lower_bound); let upper_bound_key = encode_bytes(&tracker.req_ctx.upper_bound); if buckets.span_count(&lower_bound_key, &upper_bound_key) > 1 { let mut bucket_not_match = errorpb::BucketVersionNotMatch::default(); bucket_not_match.set_version(buckets.version); bucket_not_match.set_keys(buckets.keys.clone().into()); let mut err = errorpb::Error::default(); err.set_bucket_version_not_match(bucket_not_match); return Err(Error::Region(err)); } } } // When snapshot is retrieved, deadline may exceed. tracker.on_snapshot_finished(); tracker.req_ctx.deadline.check()?; tracker.buckets = latest_buckets; let buckets_version = tracker.buckets.as_ref().map_or(0, |b| b.version); let mut handler = if can_handle_by_cached_request(&snapshot, &tracker.req_ctx) { // Build a cached request handler instead if cache version is matching. CachedRequestHandler::builder()(snapshot, &tracker.req_ctx)? } else { handler_builder(snapshot, &tracker.req_ctx)? }; Self::refund_remote_dispatch_precharge_if_local( &request_memory_quota_guard, remote_dispatch_precharge_bytes, ); tracker.on_slow_down(); let (mut request_type, mut is_remote) = (RequestType::Unknown, false); GLOBAL_TRACKERS.with_tracker(tracker_token, |tracker| { request_type = tracker.req_info.request_type; is_remote = tracker.req_info.is_remote; }); if is_remote { if let Some(recorder) = &poll_cpu_recorder { let recorder = recorder.clone(); handler.set_fallback_cpu_token_observer(Box::new(move |handle| { recorder.set_token_handle(handle); })); } } let priority = tracker.req_ctx.context.get_priority(); let legacy_dag_admission_allowed = read_limiter .as_ref() .is_some_and(ReadLimiter::is_legacy_read_limiter_allowed); let is_low_priority_dag = legacy_dag_admission_allowed && !is_remote && read_limiter.as_ref().is_some_and(|limiter| { is_low_priority_dag_request(limiter, request_type, priority) }); let deadline_remaining = tracker.req_ctx.deadline.check().unwrap_or_default(); if resource_debug { info!( "Handling unary cop request"; "keyspace_id" => tracker.req_ctx.context.get_keyspace_id(), "region_id" => tracker.req_ctx.context.region_id, "peer" => ?tracker.req_ctx.context.peer, "priority" => ?priority, "request_type" => ?request_type, "is_remote" => is_remote, "is_low_priority_dag" => is_low_priority_dag, "deadline_remaining" => ?deadline_remaining, ); } // Pre-admission: allocate CPU tokens with deadline-aware retry. // // TODO(#5480): intermediate fair-gate state. When the fair global CPU gate // is active, this legacy DAG admission path is intentionally skipped, // including the low-priority DAG sub-bucket, keyspace buckets, and tenant // caps. Admission then happens in the fair read-pool wrapper against the // global bucket only. Converge these paths when keyspace-level throttle // moves into the fair pool. let (cpu_token_handle, estimated_cpu_us_opt) = if let Some(read_limiter) = read_limiter.as_ref().filter(|_| { legacy_dag_admission_allowed && !fair_global_cpu_gate_active && !is_remote && matches!(request_type, RequestType::CoprocessorDag) }) { let (handle, estimated_cpu_us) = Self::allocate_dag_cpu_tokens( &tracker.req_ctx, read_limiter, keyspace_read_limiter.as_ref(), resource_dry_run, resource_debug, priority, ) .await?; if let (Some(recorder), Some(handle)) = (&poll_cpu_recorder, &handle) { recorder.set_token_handle(handle.clone()); } (handle, estimated_cpu_us) } else { (None, None) }; tracker.on_begin_all_items(); // Step 2: Execute request with CPU token checking and CpuObserveFuture // Extract necessary values before passing tracker to track() let deadline = tracker.req_ctx.deadline; let keyspace_id = tracker.req_ctx.context.get_keyspace_id(); let region_id = tracker.req_ctx.context.region_id; let peer = tracker.req_ctx.context.peer.clone().into_option(); let runtime_check_handle = if fair_global_cpu_gate_active || !legacy_dag_admission_allowed { None } else { prepare_cpu_token_runtime_check( handler.as_mut(), cpu_token_handle.as_ref().cloned(), resource_debug, keyspace_id, region_id, ) }; let handle_request_future = check_deadline(handler.handle_request(), deadline); let handle_request_future = track(handle_request_future, &mut tracker); let cpu_metrics_ctx = CpuThrottleMetricsContext { keyspace_id, region_id, peer, request_type, is_remote, is_low_priority_dag, estimated_cpu_us: estimated_cpu_us_opt, resource_publisher: resource_publisher.clone(), resource_debug, }; let (cpu_time, deadline_res) = run_with_cpu_token_check( handle_request_future, semaphore.as_deref(), runtime_check_handle, resource_debug, keyspace_id, region_id, ) .await?; if let Some(ref handle) = cpu_token_handle { finalize_cpu_throttle_metrics( &cpu_metrics_ctx, handle, cpu_time, poll_cpu_recorder.is_none(), ); } // Keep the handle alive so it gets dropped and tokens are returned drop(cpu_token_handle); let result = deadline_res.map_err(Error::from).and_then(|res| res); // There might be errors when handling requests. In this case, we still need its // execution metrics. let mut exec_summary = ExecSummary::default(); handler.collect_scan_summary(&mut exec_summary); tracker.collect_scan_process_time(exec_summary); let mut storage_stats = Statistics::default(); handler.collect_scan_statistics(&mut storage_stats); let total_size = storage_stats.total_size; tracker.collect_storage_statistics(storage_stats); let final_is_remote = { let mut final_is_remote = is_remote; GLOBAL_TRACKERS.with_tracker(tracker_token, |tracker| { final_is_remote = tracker.req_info.is_remote; }); final_is_remote }; if is_remote && !final_is_remote { Self::refund_remote_dispatch_precharge_if_local( &request_memory_quota_guard, remote_dispatch_precharge_bytes, ); } let (exec_details, exec_details_v2) = tracker.get_exec_details(); tracker.on_finish_all_items(); record_logical_read_bytes(exec_details_v2.get_scan_detail_v2().processed_versions_size); // Using "process_wall_time_ns"(in exec_details) here is better than // "time_processed_ns" (in ExecSummary) as the latter may also include // task suspending duration. let processed_time = Duration::from_nanos(exec_details_v2.get_time_detail_v2().process_wall_time_ns); if let Some(lazy_remote_pattern) = tracker.req_ctx.lazy_remote_pattern.clone() { if let Some(remote_ctx) = remote_ctx { remote_ctx.report_lazy_remote_pattern( lazy_remote_pattern, total_size, processed_time, ); } } LOCAL_COPR_CPU_STATS.with_borrow_mut(|local_stats| { local_stats.record_cpu_time(tracker.req_ctx.context.get_keyspace_id(), processed_time); }); let mut resp = match result { Ok(resp) => { let resp_size = resp.data.len() as u64; COPR_RESP_SIZE.inc_by(resp_size); record_network_out_bytes(resp_size); resp } Err(e) => make_error_response(e).into(), }; resp.set_exec_details(exec_details); if !resp.has_exec_details_v2() { resp.set_exec_details_v2(exec_details_v2); } // Fill RUV2 coprocessor-specific fields. let resp_bytes = resp.get_data().len() as u64; handler.set_ru_v2_metrics(resp.mut_exec_details_v2().mut_ru_v2(), resp_bytes); resp.set_latest_buckets_version(buckets_version); if let Some(resource_publisher) = &resource_publisher { if !final_is_remote && matches!(request_type, RequestType::CoprocessorDag) && total_size > 0 && read_limiter .as_ref() .is_some_and(ReadLimiter::is_legacy_read_limiter_enabled) { resource_publisher.publish(ResourceEvent::CollectMetric(Metric { resource: Resource::Read { bytes: total_size as u64, }, scope: Scope::Region { keyspace_id: tracker.req_ctx.context.get_keyspace_id(), region_id: tracker.req_ctx.context.region_id, }, severity: Severity::Normal, })); } } Ok(resp) } /// Handle a unary request and run on the read pool. /// /// Returns `Err(err)` if the read pool is full. Returns `Ok(future)` in /// other cases. The future inside may be an error however. fn handle_unary_request( &self, r: ParseCopRequestResult, ) -> impl Future>> { let remote_dispatch_precharge_bytes = r.remote_dispatch_precharge_bytes; let req_ctx = r.req_ctx; let req_tag = r.req_tag; let priority = req_ctx.context.get_priority(); let keyspace_id = req_ctx.context.get_keyspace_id(); let region_id = req_ctx.context.region_id; let task_id = req_ctx.build_task_id(); #[cfg(feature = "testexport")] let task_id = crate::read_pool::test_util::mock_spawn_task_id( req_ctx.context.get_resource_group_tag(), req_ctx.context.get_region_id(), ) .unwrap_or(task_id); let request_deadline = req_ctx.deadline.to_std_instant(); let raw_key_ranges = if is_user_read_req(req_tag) { req_ctx .ranges .iter() .map(|key_range| (key_range.get_start().to_vec(), key_range.get_end().to_vec())) .collect() } else { Vec::new() }; let resource_tag = self .resource_tag_factory .new_tag_with_key_ranges(&req_ctx.context, raw_key_ranges); let mut allocated_bytes = resource_tag.approximate_heap_size(); // box the tracker so that moving it is cheap. let tracker = Box::new(Tracker::new(req_ctx, r.req_tag, self.slow_log_threshold)); allocated_bytes += tracker.approximate_mem_size(); let remote_ctx = self.remote_ctx.clone(); let (resource_enabled, resource_dry_run, resource_debug) = self .resource_controller .as_ref() .map(|controller| { ( controller.get_enabled(), controller.get_dry_run(), controller.get_debug(), ) }) .unwrap_or_default(); let (resource_publisher, read_limiter, keyspace_read_limiter) = if resource_enabled { let read_limiter = self.read_limiter.clone(); let keyspace_read_limiter = read_limiter .as_ref() .and_then(|limiter| limiter.get_limiter(keyspace_id)); ( self.resource_publisher.clone(), read_limiter, keyspace_read_limiter, ) } else { (None, None, None) }; let tracker_token = get_tls_tracker_token(); let (tx, rx) = oneshot::channel(); let request_memory_quota_guard = Arc::new(Mutex::new(OwnedAllocated::new(self.memory_quota.clone()))); let fair_cpu_task_type = match req_tag { ReqTag::select => Some(FairCpuTaskType::DagSelect), ReqTag::index => Some(FairCpuTaskType::DagIndex), _ => None, }; let poll_cpu_recorder = read_limiter.as_ref().and_then(|read_limiter| { if !read_limiter.is_cpu_throttle_enabled() || fair_cpu_task_type.is_none() { return None; } resource_publisher.as_ref().map(|publisher| { Arc::new(CoprReadPoolPollCpuRecorder::new( keyspace_id, region_id, resource_debug, publisher.clone(), )) }) }); let fair_global_cpu_gate_active = self.read_pool.uses_fair_global_cpu_gate() && poll_cpu_recorder.is_some(); let future = Self::handle_unary_request_impl( self.semaphore.clone(), tracker, r.handler_builder, remote_dispatch_precharge_bytes, Some(Arc::clone(&request_memory_quota_guard)), remote_ctx, tracker_token, resource_publisher, poll_cpu_recorder.clone(), fair_global_cpu_gate_active, read_limiter, keyspace_read_limiter, resource_dry_run, resource_debug, ) .in_resource_metering_tag(resource_tag) .map(|res| { sync_read_pool_metrics_to_current_tracker(true); let _ = tx.send(res); }); let task_options = match (poll_cpu_recorder.as_ref(), fair_cpu_task_type) { (Some(recorder), Some(task_type)) => ReadPoolTaskOptions::default() .with_poll_cpu_recorder(recorder.clone() as Arc) .with_fair_cpu_admission(task_type, Some(request_deadline)), _ => ReadPoolTaskOptions::default(), }; let res = self.read_pool_spawn_with_memory_quota_check( allocated_bytes, remote_dispatch_precharge_bytes, future, request_memory_quota_guard, priority, keyspace_id, task_id, task_options, ); TrackedFuture::new(async move { res?; rx.map_err(|_| Error::MaxPendingTasksExceeded).await? }) } /// Parses and handles a unary request. Returns a future that will never /// fail. If there are errors during parsing or handling, they will be /// converted into a `Response` as the success result of the future. #[inline] pub fn parse_and_handle_unary_request( &self, mut req: coppb::Request, peer: Option, ) -> impl Future> { let overload_protector = self.overload_protector.clone(); let start_ts = req.start_ts; if let Some(protector) = &overload_protector { if let Some(stats) = protector.is_overloaded(start_ts) { let store_id = req.get_context().get_peer().get_store_id(); return async move { make_error_response(Error::OverloadProtection(format!( "{:?} on store {}", stats, store_id ))) .into() } .boxed(); } } { let ctx_ref = req.get_context(); let keyspace_id = ctx_ref.get_keyspace_id(); let region_id = ctx_ref.get_region_id(); let range_count = req.get_ranges().len(); let first_range = req .get_ranges() .first() .map(|r| (r.get_start().to_vec(), r.get_end().to_vec())); let last_range = req .get_ranges() .last() .map(|r| (r.get_start().to_vec(), r.get_end().to_vec())); txn_debug!(trace_event::types::Category::ReqResp, "Coprocessor request entry"; "start_ts" => req.start_ts, "keyspace_id" => keyspace_id, "region_id" => region_id, "range_count" => range_count, "first_range_start" => first_range .as_ref() .map(|r| log_wrappers::Value::key(&r.0)), "first_range_end" => first_range .as_ref() .map(|r| log_wrappers::Value::key(&r.1)), "last_range_start" => last_range .as_ref() .map(|r| log_wrappers::Value::key(&r.0)), "last_range_end" => last_range .as_ref() .map(|r| log_wrappers::Value::key(&r.1)), "request_source" => %ctx_ref.get_request_source(), "peer" => ?peer ); } let tracker = GLOBAL_TRACKERS.insert(::tracker::Tracker::new(RequestInfo::new( req.get_context(), RequestType::Unknown, req.start_ts, ))); let result_of_batch = self.process_batch_tasks(&mut req, &peer); set_tls_tracker_token(tracker); with_tls_tracker(|tracker| { tracker.metrics.grpc_req_size = req.compute_size() as u64; }); let result_of_future = self .parse_request_and_check_memory_locks( req, peer, false, self.max_resp_size.load(Ordering::Relaxed), ) .map(|r| self.handle_unary_request(r)); async move { let res = match result_of_future { Err(e) => { let mut res = make_error_response(e); let batch_res = result_of_batch.await; res.set_batch_responses(batch_res.into()); res.into() } Ok(handle_fut) => { let (handle_res, batch_res) = futures::join!(handle_fut, result_of_batch); let mut res = handle_res.unwrap_or_else(|e| make_error_response(e).into()); res.set_batch_responses(batch_res.into()); GLOBAL_TRACKERS.with_tracker(tracker, |tracker| { let exec_details_v2 = res.mut_exec_details_v2(); tracker.write_scan_detail(exec_details_v2.mut_scan_detail_v2()); if let Some(details) = tracker.read_pool_task_details() { exec_details_v2.set_read_pool_task_details(details); } }); if let Some(protector) = overload_protector { protector.report_cop_task_stats(make_cop_task_stats(start_ts, &res)); } res } }; GLOBAL_TRACKERS.remove(tracker); res } .boxed() } // process_batch_tasks process the input batched coprocessor tasks if any, // prepare all the requests and schedule them into the read pool, then // collect all the responses and convert them into the `StoreBatchResponse` // type. fn process_batch_tasks( &self, req: &mut coppb::Request, peer: &Option, ) -> impl Future> { let mut batch_futs = Vec::with_capacity(req.tasks.len()); let previous_tracker = get_tls_tracker_token(); let mut tasks = req.take_tasks(); let base_req = req.clone(); let batch_reqs: Vec<(coppb::Request, u64)> = tasks .iter_mut() .map(|task| Self::make_store_batch_request(&base_req, task)) .collect(); for (cur_req, task_id) in batch_reqs.into_iter() { let request_info = RequestInfo::new( cur_req.get_context(), RequestType::Unknown, cur_req.start_ts, ); let cur_tracker = GLOBAL_TRACKERS.insert(::tracker::Tracker::new(request_info)); set_tls_tracker_token(cur_tracker); let parse_result = self.parse_request_and_check_memory_locks( cur_req, peer.clone(), false, self.max_resp_size.load(Ordering::Relaxed), ); let mut response = coppb::StoreBatchTaskResponse::new(); response.set_task_id(task_id); match parse_result { Ok(r) => { let fut = self.handle_unary_request(r); set_tls_tracker_token(previous_tracker); let fut = async move { let res = fut.await; match res { Ok(mut resp) => { response.set_data(resp.take_data()); if let Some(err) = resp.region_error.take() { response.set_region_error(err); } if let Some(lock_info) = resp.locked.take() { response.set_locked(lock_info); } response.set_other_error(resp.take_other_error()); // keep the exec details already generated. response.set_exec_details_v2(resp.take_exec_details_v2()); GLOBAL_TRACKERS.with_tracker(cur_tracker, |tracker| { tracker.write_scan_detail( response.mut_exec_details_v2().mut_scan_detail_v2(), ); }); } Err(e) => { make_error_batch_response(&mut response, e); } } GLOBAL_TRACKERS.with_tracker(cur_tracker, |tracker| { if let Some(details) = tracker.read_pool_task_details() { response .mut_exec_details_v2() .set_read_pool_task_details(details); } }); GLOBAL_TRACKERS.remove(cur_tracker); response }; batch_futs.push(future::Either::Left(fut)); } Err(e) => { set_tls_tracker_token(previous_tracker); GLOBAL_TRACKERS.remove(cur_tracker); batch_futs.push(future::Either::Right(async move { make_error_batch_response(&mut response, e); response })); } } } stream::FuturesOrdered::from_iter(batch_futs).collect() } fn make_store_batch_request( req: &coppb::Request, task: &mut coppb::StoreBatchTask, ) -> (coppb::Request, u64) { let mut new_req = req.clone(); // Disable the coprocessor cache path for the batched tasks, the // coprocessor cache related fields are not passed in the "task" by now. new_req.is_cache_enabled = false; new_req.ranges = task.take_ranges(); let new_context = new_req.mut_context(); new_context.set_region_id(task.get_region_id()); new_context.set_region_epoch(task.take_region_epoch()); new_context.set_peer(task.take_peer()); let buckets_version = task.get_buckets_version(); if buckets_version != 0 { new_context.set_buckets_version(buckets_version); } (new_req, task.get_task_id()) } /// The real implementation of handling a stream request. /// /// It first retrieves a snapshot, then builds the `RequestHandler` over the /// snapshot and the given `handler_builder`. Finally, it calls the stream /// request interface of the `RequestHandler` multiple times to process the /// request and produce multiple results. fn handle_stream_request_impl( semaphore: Option>, mut tracker: Box>, handler_builder: RequestHandlerBuilder, remote_dispatch_precharge_bytes: usize, request_memory_quota_guard: Option, ) -> impl futures::stream::Stream> { try_stream! { let _permit = if let Some(semaphore) = semaphore.as_ref() { Some(semaphore.acquire().await.expect("the semaphore never be closed")) } else { None }; // When this function is being executed, it may be queued for a long time, so that // deadline may exceed. tracker.on_scheduled(); tracker.req_ctx.deadline.check()?; // Safety: spawning this function using a `FuturePool` ensures that a TLS engine // exists. let snapshot = unsafe { with_tls_engine(|engine| Self::async_snapshot(engine, &tracker.req_ctx)) } .await?; // When snapshot is retrieved, deadline may exceed. tracker.on_snapshot_finished(); tracker.req_ctx.deadline.check()?; let mut handler = handler_builder(snapshot, &tracker.req_ctx)?; Self::refund_remote_dispatch_precharge_if_local( &request_memory_quota_guard, remote_dispatch_precharge_bytes, ); tracker.on_slow_down(); tracker.on_begin_all_items(); loop { let result = { tracker.on_begin_item(); let result = handler.handle_streaming_request().await; let mut storage_stats = Statistics::default(); handler.collect_scan_statistics(&mut storage_stats); tracker.on_finish_item(Some(storage_stats)); result }; let (exec_details, exec_details_v2) = tracker.get_item_exec_details(); LOCAL_COPR_CPU_STATS.with_borrow_mut(|local_stats| { local_stats.record_cpu_time(tracker.req_ctx.context.get_keyspace_id(), Duration::from_nanos(exec_details_v2.get_time_detail_v2().process_wall_time_ns)); }); match result { Err(e) => { let mut resp = make_error_response(e); resp.set_exec_details(exec_details); resp.set_exec_details_v2(exec_details_v2); let resp_bytes = resp.get_data().len() as u64; handler.set_ru_v2_metrics(resp.mut_exec_details_v2().mut_ru_v2(), resp_bytes); yield resp; break; }, Ok((None, _)) => break, Ok((Some(mut resp), finished)) => { let resp_size = resp.data.len() as u64; COPR_RESP_SIZE.inc_by(resp_size); record_network_out_bytes(resp_size); record_logical_read_bytes(exec_details_v2.get_scan_detail_v2().processed_versions_size); resp.set_exec_details(exec_details); resp.set_exec_details_v2(exec_details_v2); let resp_bytes = resp.get_data().len() as u64; handler.set_ru_v2_metrics(resp.mut_exec_details_v2().mut_ru_v2(), resp_bytes); yield resp; if finished { break; } } } } tracker.on_finish_all_items(); } } /// Handle a stream request and run on the read pool. /// /// Returns `Err(err)` if the read pool is full. Returns `Ok(stream)` in /// other cases. The stream inside may produce errors however. fn handle_stream_request( &self, r: ParseCopRequestResult, ) -> Result>> { let remote_dispatch_precharge_bytes = r.remote_dispatch_precharge_bytes; let req_tag = r.req_tag; let req_ctx = r.req_ctx; let (tx, rx) = mpsc::channel::>(self.stream_channel_size); let priority = req_ctx.context.get_priority(); let keyspace_id = req_ctx.context.get_keyspace_id(); let raw_key_ranges = if is_user_read_req(req_tag) { req_ctx .ranges .iter() .map(|key_range| (key_range.get_start().to_vec(), key_range.get_end().to_vec())) .collect() } else { Vec::new() }; let resource_tag = self .resource_tag_factory .new_tag_with_key_ranges(&req_ctx.context, raw_key_ranges); let mut allocated_bytes = resource_tag.approximate_heap_size(); let task_id = req_ctx.build_task_id(); let tracker = Box::new(Tracker::new(req_ctx, r.req_tag, self.slow_log_threshold)); allocated_bytes += tracker.approximate_mem_size(); let request_memory_quota_guard = Arc::new(Mutex::new(OwnedAllocated::new(self.memory_quota.clone()))); let future = Self::handle_stream_request_impl( self.semaphore.clone(), tracker, r.handler_builder, remote_dispatch_precharge_bytes, Some(Arc::clone(&request_memory_quota_guard)), ) .in_resource_metering_tag(resource_tag) .then(futures::future::ok::<_, mpsc::SendError>) .forward(tx) .unwrap_or_else(|e| { warn!("coprocessor stream send error"; "error" => %e); }); self.read_pool_spawn_with_memory_quota_check( allocated_bytes, remote_dispatch_precharge_bytes, future, request_memory_quota_guard, priority, keyspace_id, task_id, ReadPoolTaskOptions::default(), )?; Ok(rx) } /// Parses and handles a stream request. Returns a stream that produce each /// result in a `Response` and will never fail. If there are errors during /// parsing or handling, they will be converted into a `Response` as the /// only stream item. #[inline] pub fn parse_and_handle_stream_request( &self, req: coppb::Request, peer: Option, ) -> impl futures::stream::Stream { let max_resp_size = if req.get_tp() == REQ_TYPE_DB9_DAG { self.max_resp_size.load(Ordering::Relaxed) } else { u64::MAX }; let result_of_stream = self .parse_request_and_check_memory_locks(req, peer, true, max_resp_size) .and_then(|r| self.handle_stream_request(r)); // Result, Error> futures::stream::once(futures::future::ready(result_of_stream)) .try_flatten() .or_else(|e| futures::future::ok(make_error_response(e))) .map(|item: std::result::Result<_, ()>| item.unwrap()) } pub fn handle_delegate_request( &self, mut req: coppb::DelegateRequest, ) -> impl Future { let mut ranges = Vec::new(); let key_ranges = req.take_ranges().into_vec(); for key_range in key_ranges.clone() { let kv_range = ( bytes::Bytes::from(key_range.start.clone()), bytes::Bytes::from(key_range.end.clone()), ); ranges.push(kv_range) } let req_ctx = ReqContext::new( req.take_context(), key_ranges, self.max_handle_duration, None, None, req.start_ts.into(), None, None, ); let check_mem_res = Endpoint::::check_memory_locks(&self.concurrency_manager, &req_ctx); let priority = req_ctx.context.get_priority(); let keyspace_id = req_ctx.context.get_keyspace_id(); let task_id = req_ctx.build_task_id(); // Legacy path from // https://github.com/tidbcloud/cloud-storage-engine/pull/698, // superseded by the HTTP-based `RemoteDagDispatcher` in // https://github.com/tidbcloud/cloud-storage-engine/pull/1096. No // production traffic reaches it today, but this route still bypasses // `read_pool_spawn_with_memory_quota_check`, so we make any unexpected // traffic explicit via the warning log and // `COPR_DELEGATE_LEGACY_REQUEST_TOTAL`. COPR_DELEGATE_LEGACY_REQUEST_TOTAL.inc(); warn!( "legacy delegate coprocessor request bypasses coprocessor memory quota admission"; "start_ts" => req.start_ts, ); let res = self .read_pool .spawn_handle( async move { check_mem_res?; let snapshot = unsafe { with_tls_engine(|engine| Self::async_snapshot(engine, &req_ctx)) } .await?; let snap_access = snapshot.get_kvengine_snap().unwrap().clone(); let cloud_store = CloudStore::new( snapshot, req.start_ts, req_ctx.bypass_locks.clone(), !req_ctx.context.get_not_fill_cache(), false, ); cloud_store.check_locks_in_range( &req_ctx.lower_bound, &req_ctx.upper_bound, false, )?; let mut resp = coppb::DelegateResponse::default(); let (mem_data, _) = snap_access.build_mem_data(&ranges, req.start_ts); resp.set_mem_table_data(mem_data); let (_, snapshot) = snap_access.marshal(ranges.as_slice(), true, false); resp.set_snapshot(snapshot); Ok(resp) }, priority, keyspace_id, task_id, ) .map_err(|_| Error::MaxPendingTasksExceeded); async move { match res.await { Ok(Ok(resp)) => resp, Ok(Err(e)) => make_error_delegate_response(e), Err(e) => make_error_delegate_response(e), } } } fn read_pool_spawn_with_memory_quota_check( &self, allocated_bytes: usize, remote_dispatch_precharge_bytes: usize, future: F, request_memory_quota_guard: SharedOwnedAllocated, priority: kvrpcpb::CommandPri, keyspace_id: u32, task_id: u64, task_options: ReadPoolTaskOptions, ) -> Result<()> where F: Future + Send + 'static, { let request_memory_quota_bytes = allocated_bytes + remote_dispatch_precharge_bytes + mem::size_of_val(&future); request_memory_quota_guard .lock() .unwrap() .alloc(request_memory_quota_bytes)?; let fut = future.map(move |_| { // Release the remaining quota after handle completed. Local requests // may already have refunded the remote-dispatch extra portion. drop(request_memory_quota_guard); }); self.read_pool .spawn_with_options(fut, priority, keyspace_id, task_id, task_options) .map_err(|_| Error::MaxPendingTasksExceeded) } fn refund_remote_dispatch_precharge_if_local( request_memory_quota_guard: &Option, remote_dispatch_precharge_bytes: usize, ) { if remote_dispatch_precharge_bytes == 0 { return; } let mut is_remote = false; with_tls_tracker(|tracker| { is_remote = tracker.req_info.is_remote; }); if !is_remote { if let Some(request_memory_quota_guard) = request_memory_quota_guard { request_memory_quota_guard .lock() .unwrap() .free(remote_dispatch_precharge_bytes); } } } } fn check_memory_locks_for_ranges( concurrency_manager: &ConcurrencyManager, req_ctx: &ReqContext, key_ranges: &[coppb::KeyRange], ) -> Result<()> { let start_ts = req_ctx.txn_start_ts; if !req_ctx.context.get_stale_read() { concurrency_manager.update_max_ts(start_ts, || format!("coprocessor-{}", start_ts))?; } if need_check_locks(req_ctx.context.get_isolation_level()) { let begin_instant = Instant::now(); for range in key_ranges { let start_key = txn_types::Key::from_raw_maybe_unbounded(range.get_start()); let end_key = txn_types::Key::from_raw_maybe_unbounded(range.get_end()); concurrency_manager .read_range_check(start_key.as_ref(), end_key.as_ref(), |key, lock| { Lock::check_ts_conflict( Cow::Borrowed(lock), key, start_ts, &req_ctx.bypass_locks, req_ctx.context.get_isolation_level(), ) }) .map_err(|e| { MEM_LOCK_CHECK_HISTOGRAM_VEC_STATIC .locked .observe(begin_instant.saturating_elapsed().as_secs_f64()); MvccError::from(e) })?; } MEM_LOCK_CHECK_HISTOGRAM_VEC_STATIC .unlocked .observe(begin_instant.saturating_elapsed().as_secs_f64()); } Ok(()) } /// Fixed, low-cardinality execution kind of a completed remote coprocessor /// request. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RemoteCopRequestKind { Select, Index, Db9Select, Db9Index, AnalyzeTable, AnalyzeIndex, AnalyzeFullSampling, ChecksumTable, ChecksumIndex, } /// Scan work completed by a remote coprocessor request before RU discounting. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct RemoteCopScanStats { pub request_kind: RemoteCopRequestKind, pub processed_keys: u64, pub total_keys: u64, pub processed_bytes: u64, pub total_bytes: u64, } /// Successful remote coprocessor response and its undiscounted scan work. pub struct RemoteCopResult { pub response: MemoryTraceGuard, pub scan_stats: RemoteCopScanStats, } pub async fn parse_request_and_handle_remote_cop( req: coppb::Request, peer: Option, max_handle_duration: Duration, max_resp_size: u64, snap: S, ) -> Result { debug_assert_ne!(get_tls_tracker_token(), INVALID_TRACKER_TOKEN); let api_version = req.get_context().get_api_version(); dispatch_api_version!( api_version, { parse_request_and_handle_remote_cop_impl::( req, peer, max_handle_duration, max_resp_size, snap, ) .await }, Err(box_err!( "API V3 is not supported on the cloud-engine branch" )) ) } pub async fn parse_request_and_handle_remote_cop_impl( req: coppb::Request, peer: Option, max_handle_duration: Duration, max_resp_size: u64, snap: S, ) -> Result { parse_request_and_handle_remote_cop_impl_with_quota::( req, peer, max_handle_duration, max_resp_size, None, snap, ) .await } async fn parse_request_and_handle_remote_cop_impl_with_quota( mut req: coppb::Request, peer: Option, max_handle_duration: Duration, max_resp_size: u64, quota_limiter: Option>, snap: S, ) -> Result { let (context, data, ranges, mut start_ts) = ( req.take_context(), req.take_data(), req.take_ranges().to_vec(), req.get_start_ts(), ); let is_cache_enabled = req.get_is_cache_enabled(); let request_type = req.get_tp(); let req_ctx: ReqContext; let request_kind: RemoteCopRequestKind; let mut handler = match request_type { REQ_TYPE_DAG => { let mut dag = DagRequest::default(); box_try!(dag.merge_from_bytes(&data)); let mut table_scan = false; let mut is_desc_scan = false; if let Some(scan) = dag.get_executors().iter().next() { table_scan = scan.get_tp() == ExecType::TypeTableScan; if table_scan { is_desc_scan = scan.get_tbl_scan().get_desc(); } else { is_desc_scan = scan.get_idx_scan().get_desc(); } } request_kind = if table_scan { RemoteCopRequestKind::Select } else { RemoteCopRequestKind::Index }; if start_ts == 0 { start_ts = dag.get_start_ts_fallback(); } req_ctx = { let mut inner = ReqContextInner::new( context, ranges, max_handle_duration, peer, Some(is_desc_scan), start_ts.into(), // Do not set `cache_matched_version` as the local endpoint // has already checked `CachedRequestHandler` eligibility // before it offloads to a remote handler. None, None, ); // FIXME: Fix the `Locked` error of async commit. inner.bypass_locks = TsSet::All; inner.into() }; with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorDag; tracker.req_info.start_ts = start_ts; }); // Remote cop receives the rebuilt request after local cache-hit // handling. Preserve cacheability checking for cache-enabled DAGs // so refreshed responses are marked cacheable only when safe. let check_has_newer_ts_data = is_cache_enabled; let store = CloudStore::new( snap, start_ts, req_ctx.bypass_locks.clone(), !req_ctx.context.get_not_fill_cache(), check_has_newer_ts_data, ); let paging_size = match req.get_paging_size() { 0 => None, i => Some(i), }; let paging_size_bytes = match req.get_paging_size_bytes() { 0 => None, i => Some(i), }; dag::DagHandlerBuilder::new( dag, req_ctx.ranges.clone(), store, tidb_query_common::storage::StubAccessor::none(), req_ctx.deadline, 64, is_cache_enabled, paging_size, paging_size_bytes, max_resp_size, ReadRescheduleOptions::default().range_scan(), ) .quota_limiter(quota_limiter) .build::()? } REQ_TYPE_DB9_DAG => { let db9 = { let mut input = CodedInputStream::from_bytes(&data); input.set_recursion_limit(1000); let mut db9 = kvproto::db9_coprocessor::Db9DagRequest::default(); if let Err(err) = db9.merge_from(&mut input) { let err = Error::Other(err.to_string()); return Err(err); } db9 }; request_kind = if is_db9_index_scan_kind(db9.get_scan().get_kind()) { RemoteCopRequestKind::Db9Index } else { RemoteCopRequestKind::Db9Select }; if let Some(reason) = unsupported_remote_db9_request_reason(&db9) { warn!( "remote DB9 cop worker rejected unsupported request"; "reason" => reason, "keyspace_id" => context.get_keyspace_id(), "start_ts" => start_ts, "scan_kind" => ?db9.get_scan().get_kind(), "has_limit" => db9.has_limit(), "ranges" => ranges.len(), ); return Err(Error::Other(reason.to_owned())); } let mut req_ctx_inner = ReqContextInner::new( context, ranges, max_handle_duration, peer, Some(db9.get_scan().get_desc()), start_ts.into(), // DB9 does not populate the normal DAG cache metadata, so // it must not use the generic request-cache fast path. None, None, ); req_ctx_inner.bypass_locks = TsSet::All; req_ctx = req_ctx_inner.into(); with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorDag; tracker.req_info.start_ts = start_ts; }); Box::new(db9::Db9DagHandler::new( db9, snap, &req_ctx, max_resp_size, tidb_query_common::storage::StubAccessor::none(), ReadRescheduleOptions::default(), )?) } REQ_TYPE_ANALYZE => { let analyze = { let mut input = CodedInputStream::from_bytes(&data); input.set_recursion_limit(1000); let mut analyze = AnalyzeReq::default(); match analyze.merge_from(&mut input) { Ok(_) => (), Err(e) => return Err(Error::Other(e.to_string())), }; analyze }; request_kind = match analyze.get_tp() { AnalyzeType::TypeIndex | AnalyzeType::TypeCommonHandle | AnalyzeType::TypeSampleIndex => RemoteCopRequestKind::AnalyzeIndex, AnalyzeType::TypeColumn | AnalyzeType::TypeMixed => { RemoteCopRequestKind::AnalyzeTable } AnalyzeType::TypeFullSampling => RemoteCopRequestKind::AnalyzeFullSampling, }; if start_ts == 0 { start_ts = analyze.get_start_ts_fallback(); } req_ctx = ReqContext::new( context, ranges, max_handle_duration, peer, None, start_ts.into(), None, None, ); with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorAnalyze; tracker.req_info.start_ts = start_ts; }); Box::new( statistics::analyze::AnalyzeContext::<_, F>::new_with_quota( analyze, req_ctx.ranges.clone(), start_ts, snap, &req_ctx, quota_limiter, None, RemoteRequest::default(), ReadRescheduleOptions::default(), ) .unwrap(), ) } REQ_TYPE_CHECKSUM => { let checksum = { let mut input = CodedInputStream::from_bytes(&data); input.set_recursion_limit(1000); let mut checksum = ChecksumRequest::default(); box_try!(checksum.merge_from(&mut input)); checksum }; request_kind = if checksum.get_scan_on() == ChecksumScanOn::Table { RemoteCopRequestKind::ChecksumTable } else { RemoteCopRequestKind::ChecksumIndex }; if start_ts == 0 { start_ts = checksum.get_start_ts_fallback(); } req_ctx = ReqContext::new( context, ranges, max_handle_duration, peer, None, start_ts.into(), None, None, ); with_tls_tracker(|tracker| { tracker.req_info.request_type = RequestType::CoprocessorChecksum; tracker.req_info.start_ts = start_ts; }); Box::new( checksum::ChecksumContext::new( checksum, req_ctx.ranges.clone(), start_ts, snap, &req_ctx, None, RemoteRequest::default(), ReadRescheduleOptions::default(), ) .unwrap(), ) } tp => return Err(Error::Other(format!("unsupported tp {}", tp))), }; let deadline = Deadline::from_now(max_handle_duration); let mut resp = check_deadline(handler.handle_request(), deadline).await??; let mut exec_summary = ExecSummary::default(); handler.collect_scan_summary(&mut exec_summary); let mut time_detail = TimeDetail::default(); time_detail.set_process_wall_time_ms((exec_summary.time_processed_ns / 1000000) as u64); let mut time_detail_v2 = TimeDetailV2::default(); time_detail_v2.set_process_wall_time_ns(exec_summary.time_processed_ns as u64); let mut storage_stats = Statistics::default(); handler.collect_scan_statistics(&mut storage_stats); let scan_stats = RemoteCopScanStats { request_kind, processed_keys: storage_stats.write.processed_keys as u64, total_keys: storage_stats.write.total_op_count() as u64, processed_bytes: storage_stats.processed_size as u64, total_bytes: storage_stats.total_size as u64, }; let mut detail_v2 = ScanDetailV2::default(); storage_stats.write_scan_detail(&mut detail_v2); with_tls_tracker(|tracker| tracker.write_scan_detail(&mut detail_v2)); // To be compatible with TiDB, we need to discount both processed_size and // total_size. // See https://github.com/tidbcloud/client-go-cse/blob/a6c924350/internal/resourcecontrol/resource_control.go#L154 let discounted_processed_size = storage_stats.processed_size * SIZE_DISCOUNT_PERCENT / 100; let discounted_total_size = storage_stats.total_size * SIZE_DISCOUNT_PERCENT / 100; detail_v2.set_processed_versions_size(discounted_processed_size as u64); detail_v2.set_total_versions_size(discounted_total_size as u64); let mut exec_details_v2 = kvrpcpb::ExecDetailsV2::default(); exec_details_v2.set_scan_detail_v2(detail_v2); exec_details_v2.set_time_detail(time_detail); exec_details_v2.set_time_detail_v2(time_detail_v2); // Fill RUV2 for remote coprocessor. handler.set_ru_v2_metrics(exec_details_v2.mut_ru_v2(), resp.get_data().len() as u64); with_tls_tracker(|tracker| { tracker.write_ru_v2(exec_details_v2.mut_ru_v2()); if let Some(details) = tracker.read_pool_task_details() { exec_details_v2.set_read_pool_task_details(details); } }); resp.set_exec_details_v2(exec_details_v2); Ok(RemoteCopResult { response: resp, scan_stats, }) } pub async fn prefetch_ia_remote_segments( tag: &str, snap_ctx: &SnapCtx, snap_access: &SnapAccess, ranges: &[coppb::KeyRange], deadline: Deadline, ) -> Result /* cache_hit */> { if snap_access.is_sync() { return Ok(None); } let IaCtx::Enabled(ia_mgr, _) = &snap_ctx.ia_ctx else { debug_assert!(false, "IA must be enabled when snap is async"); return Ok(None); }; fail_point!("prefetch_ia_remote_segments_err", |action| { match action.as_deref() { Some("deadline") => Err(Error::DeadlineExceeded), _ => Err(box_err!("injected prefetch segments failure")), } }); let mut segments = vec![]; let mut total_segments = 0; for ran in ranges { let (segs, total) = snap_access .get_ia_remote_segments((ran.get_start(), ran.get_end())) .map_err(|err| -> Error { error!("{} get remote segments failed", tag; "range" => ?ran, "err" => ?err); box_err!("get remote segments failed: {:?}", err) })?; segments.extend(segs); total_segments += total; } if segments.is_empty() { return Ok(Some(1.0)); } segments.sort(); segments.dedup(); info!("{} prefetch remote segments", tag; "segments" => ?segments); let prefetch_segments = segments.len(); let _enter = ia_mgr.enter_runtime(); // Spawn tasks on runtime of `ia_mgr`. let keyspace_id = snap_access.get_keyspace_id(); let from_ia_region = snap_access.from_ia_region(); let mut tasks = vec![]; for (ident, ftype) in segments { let tag = tag.to_string(); let mgr = ia_mgr.clone(); let dfs = snap_ctx.dfs.clone(); let task = async move { mgr.prefetch_segment( ident, ftype, keyspace_id, deadline, Some(dfs.deref()), from_ia_region, ).map_err(|err| -> Error { error!("{} prefetch segment failed", tag; "ident" => %ident, "ftype" => ?ftype, "err" => ?err); if let kvengine::table::Error::DeadlineExceeded(_) = err { Error::DeadlineExceeded } else { box_err!("prefetch segment failed: {}: {:?}", ident, err) } }).await }; tasks.push(tokio::spawn(TrackedFuture::new(task))); } drop(_enter); match tokio::time::timeout_at(deadline.to_tokio_instant(), try_join_all(tasks)).await { Ok(res) => { let res = res.map_err(|err| -> Error { box_err!("{} prefetch: join tasks failed: {:?}", tag, err) })?; res.into_iter().collect::>()?; let cache_hit = (total_segments - prefetch_segments) as f64 / total_segments as f64; info!("{} prefetch ia remote segments", tag; "prefetch_segments" => prefetch_segments, "cache_hit" => cache_hit); Ok(Some(cache_hit)) } Err(_) => Err(Error::DeadlineExceeded), } } macro_rules! make_error_response_common { ($resp:expr, $tag:expr, $e:expr) => {{ match $e { Error::Region(e) => { $tag = storage::get_tag_from_header(&e); $resp.set_region_error(e); } Error::Locked(info) => { $tag = "meet_lock"; $resp.set_locked(info); } Error::DeadlineExceeded => { $tag = "deadline_exceeded"; let mut err = errorpb::Error::default(); set_deadline_exceeded_busy_error(&mut err); err.set_message($e.to_string()); $resp.set_region_error(err); } Error::MaxPendingTasksExceeded => { $tag = "max_pending_tasks_exceeded"; let reason = $e.to_string(); let mut errorpb = errorpb::Error::default(); errorpb.set_message(reason.clone()); errorpb.set_server_is_busy(new_server_is_busy( ServerIsBusyType::CoprocessorMaxPendingTasks, reason, )); $resp.set_region_error(errorpb); } Error::OverloadProtection(_) => { $tag = "overload_protection"; $resp.set_other_error($e.to_string()); } Error::RemoteNetwork(_) => { $tag = "remote_network"; $resp.set_other_error($e.to_string()); } Error::RemoteServiceUnavailable(_) => { $tag = "remote_service_unavailable"; $resp.set_other_error($e.to_string()); } Error::InvalidMaxTsUpdate(e) => { $tag = "invalid_max_ts_update"; let mut err = errorpb::Error::default(); err.set_message(e.to_string()); $resp.set_region_error(err); } Error::CpuThrottled(_) => { $tag = "cpu_throttled"; let reason = $e.to_string(); let mut errorpb = errorpb::Error::default(); errorpb.set_message(reason.clone()); errorpb.set_server_is_busy(new_server_is_busy( ServerIsBusyType::CoprocessorCpuThrottled, reason, )); $resp.set_region_error(errorpb); } Error::MemoryQuotaExceeded => { $tag = "memory_quota_exceeded"; let reason = $e.to_string(); let mut errorpb = errorpb::Error::default(); errorpb.set_message(reason.clone()); errorpb.set_server_is_busy(new_server_is_busy( ServerIsBusyType::CoprocessorMemoryQuota, reason, )); $resp.set_region_error(errorpb); } Error::Other(_) => { $tag = "other"; warn!("unexpected other error encountered processing coprocessor task"; "error" => ?&$e, ); $resp.set_other_error($e.to_string()); } }; COPR_REQ_ERROR.with_label_values(&[$tag]).inc(); }}; } fn make_error_batch_response(batch_resp: &mut coppb::StoreBatchTaskResponse, e: Error) { debug!( "batch cop task error-response"; "err" => %e ); let tag; make_error_response_common!(batch_resp, tag, e); } pub fn make_error_response(e: Error) -> coppb::Response { debug!( "error-response"; "err" => %e ); let tag; let mut resp = coppb::Response::default(); make_error_response_common!(resp, tag, e); resp } fn make_error_delegate_response(e: Error) -> coppb::DelegateResponse { warn!( "error-response"; "err" => %e ); let mut resp = coppb::DelegateResponse::default(); let tag; match e { Error::Region(e) => { tag = storage::get_tag_from_header(&e); resp.set_region_error(e); } Error::Locked(info) => { tag = "meet_lock"; resp.set_locked(info); } Error::DeadlineExceeded => { tag = "deadline_exceeded"; resp.set_other_error(e.to_string()); } Error::MaxPendingTasksExceeded => { tag = "max_pending_tasks_exceeded"; let reason = e.to_string(); let mut errorpb = errorpb::Error::default(); errorpb.set_message(reason.clone()); errorpb.set_server_is_busy(new_server_is_busy( ServerIsBusyType::CoprocessorMaxPendingTasks, reason, )); resp.set_region_error(errorpb); } Error::OverloadProtection(_) => { tag = "overload_protection"; resp.set_other_error(e.to_string()); } Error::RemoteNetwork(_) => { tag = "remote_network"; resp.set_other_error(e.to_string()); } Error::RemoteServiceUnavailable(_) => { tag = "remote_service_unavailable"; resp.set_other_error(e.to_string()); } Error::CpuThrottled(_) => { tag = "cpu_throttled"; let reason = e.to_string(); let mut errorpb = errorpb::Error::default(); errorpb.set_message(reason.clone()); errorpb.set_server_is_busy(new_server_is_busy( ServerIsBusyType::CoprocessorCpuThrottled, reason, )); resp.set_region_error(errorpb); } Error::Other(_) => { tag = "other"; resp.set_other_error(e.to_string()); } Error::InvalidMaxTsUpdate(_) => { tag = "invalid_max_ts_update"; resp.set_other_error(e.to_string()); } Error::MemoryQuotaExceeded => { tag = "memory_quota_exceeded"; let reason = e.to_string(); let mut errorpb = errorpb::Error::default(); errorpb.set_message(reason.clone()); errorpb.set_server_is_busy(new_server_is_busy( ServerIsBusyType::CoprocessorMemoryQuota, reason, )); resp.set_region_error(errorpb); } }; COPR_REQ_ERROR.with_label_values(&[tag]).inc(); resp } fn make_cop_task_stats(start_ts: u64, resp: &coppb::Response) -> CopTaskStats { let mut stats = CopTaskStats { task_id: start_ts, num_reqs: 1, ..Default::default() }; let exec_details = resp.get_exec_details_v2(); stats.num_keys = exec_details.get_scan_detail_v2().get_processed_versions(); stats.data_size = exec_details .get_scan_detail_v2() .get_processed_versions_size(); stats.duration_ms = exec_details.get_time_detail().get_process_wall_time_ms() as u32; for batch_resp in resp.get_batch_responses() { stats.num_reqs += 1; let exec_details = batch_resp.get_exec_details_v2(); stats.num_keys += exec_details.get_scan_detail_v2().get_processed_versions(); stats.data_size += exec_details .get_scan_detail_v2() .get_processed_versions_size(); stats.duration_ms += exec_details.get_time_detail().get_process_wall_time_ms() as u32; } stats } /// ExtraSnapStoreAccessor is used to get the snapshot stores for the /// extra regions. /// The "extra region" means the regions that are not the source region in a /// request. /// For example, if a cop-task contains a `IndexLookUp` executor which needs to /// access look up the primary rows, it will use this accessor to locate and get /// the snapshot of the regions which these primary rows located. #[derive(Clone)] pub struct ExtraSnapStoreAccessor { region_info_accessor: RegionInfoAccessor, concurrency_manager: ConcurrencyManager, store_id: u64, req_ctx: ReqContext, read_reschedule_options: ReadRescheduleOptions, _phantom: PhantomData E>, } impl ExtraSnapStoreAccessor { /// Creates an accessor using a request-stable read-reschedule /// configuration snapshot. Unsupported requests return `None` and access /// only the source region. pub fn new( req_ctx: ReqContext, region_info_accessor: Option, concurrency_manager: ConcurrencyManager, read_reschedule_options: ReadRescheduleOptions, ) -> Option { match region_info_accessor { Some(region_info_accessor) => { let pb_ctx = &req_ctx.context; let store_id = match pb_ctx.peer.as_ref() { // Though it is possible that the request carries a wrong store_id, we can still // use it to find the peer in the region because it will be // validated in get snapshot phase. Some(peer) => peer.get_store_id(), None => return None, }; if pb_ctx.get_isolation_level() == kvrpcpb::IsolationLevel::Si && !pb_ctx.get_stale_read() && !pb_ctx.get_replica_read() { // Some scenes are not supported before an effective argument, including: // - stale-read & replica-read, TODO: support it later // - non-SI isolation level, TODO: support it later return Some(Self { region_info_accessor, concurrency_manager, store_id, req_ctx, read_reschedule_options, _phantom: PhantomData, }); } None } None => None, } } #[inline] fn err(s: S) -> StorageError { StorageError::from(anyhow!("{}", s)) } } #[async_trait] impl RegionStorageAccessor for ExtraSnapStoreAccessor { type Storage = CloudStore; /// find the region by the specified key. /// The argument `key` should be the comparable format, you should use /// `Key::from_raw` encode the raw key. async fn find_region_by_key(&self, key: &[u8]) -> FindRegionResult { let key_in_vec = key.to_vec(); match self.region_info_accessor.find_region_info_by_key(key) { Some(info) => { if info.region.get_start_key() <= key_in_vec.as_slice() { return FindRegionResult::with_found(info.region.clone(), info.role); } FindRegionResult::with_not_found(Some(info.region.start_key.clone())) } None => FindRegionResult::with_not_found(None), } } async fn get_local_region_storage( &self, region: &metapb::Region, key_range: &[coppb::KeyRange], ) -> StorageResult { let peer = match find_peer(region, self.store_id) { Some(peer) => peer.clone(), None => { return Err(Self::err(format!( "cannot find peer in region, region_id: {}, request store_id: {}", self.store_id, region.get_id() ))); } }; check_memory_locks_for_ranges(&self.concurrency_manager, &self.req_ctx, key_range)?; let start_ts = self.req_ctx.txn_start_ts; let extra = ExtraRegionOverride { region_id: region.get_id(), region_epoch: region.get_region_epoch().clone(), peer, // Do not need to check the term because only readonly requests are // supported currently. check_term: None, }; let snapshot = unsafe { with_tls_engine(|engine: &mut E| { Endpoint::async_snapshot_for_extra(engine, &self.req_ctx, extra) }) } .await?; if snapshot.get_kvengine_snap().is_none() { return Err(Self::err("unexpected snapshot")); } tikv_util::set_current_region(region.get_id()); Ok(CloudStore::new( snapshot, start_ts.into_inner(), self.req_ctx.bypass_locks.clone(), !self.req_ctx.context.get_not_fill_cache(), // If `check_has_newer_ts_data` is enabled, the internal scanner will check if there is // newer data in the region to tell the TiDB whether to cache the cop response or not. // However, the current cop-cache implementation only supports the source region. // So here we force to set it to `false` to avoid caching the cop response when // any other region snapshot is accessed in a request. false, ) .with_read_reschedule_options(self.read_reschedule_options)) } fn get_original_region_id(&self) -> Option { Some(self.req_ctx.context.get_region_id()) } } struct KeyspaceCpuCounter { cpu_us: u64, last_updated_at: Instant, } struct ThreadLocalKeyspaceCpuStats { last_flush_time: Instant, cpu_metrics: HashMap, } impl ThreadLocalKeyspaceCpuStats { fn new() -> Self { Self { last_flush_time: Instant::now_coarse(), cpu_metrics: HashMap::new(), } } const EVICT_LOCAL_METRICS_IDLE_DUR: Duration = Duration::from_secs(600); const FLUSH_METRICS_DUR: Duration = Duration::from_secs(1); fn record_cpu_time(&mut self, keyspace_id: u32, dur: Duration) { let now = Instant::now_coarse(); let cpu_us = duration_to_us(dur); if let Some(entry) = self.cpu_metrics.get_mut(&keyspace_id) { entry.cpu_us += cpu_us; entry.last_updated_at = now; } else { self.cpu_metrics.insert( keyspace_id, KeyspaceCpuCounter { cpu_us, last_updated_at: now, }, ); } self.maybe_flush_metrics(now); } #[inline] fn maybe_flush_metrics(&mut self, now: Instant) { if now.saturating_duration_since(self.last_flush_time) >= Self::FLUSH_METRICS_DUR { let mut usages = Vec::with_capacity(self.cpu_metrics.len()); self.cpu_metrics.retain(|keyspace_id, metric| { if metric.cpu_us > 0 { usages.push((*keyspace_id, metric.cpu_us)); metric.cpu_us = 0; } now.saturating_duration_since(metric.last_updated_at) <= Self::EVICT_LOCAL_METRICS_IDLE_DUR }); report_keyspace_cpu_time(usages); self.cpu_metrics.shrink_to_fit(); self.last_flush_time = now; } } } #[cfg(test)] mod tests { use std::{ sync::{ Mutex, atomic::{self, AtomicBool}, mpsc, }, thread, vec, }; use ::tracker::{ GLOBAL_TRACKERS, RequestInfo, RequestType, clear_tls_tracker_token, set_tls_tracker_token, }; use cloud_server::mock_kv_engine::TestKvEngine; use futures::executor::{block_on, block_on_stream}; use kvproto::{ db9_coprocessor::{ Db9BinaryExpr, Db9BinaryOp, Db9ColumnInfo, Db9ColumnRefExpr, Db9ConstantExpr, Db9DagRequest, Db9Expr, Db9FuncCallExpr, Db9IndexInfo, Db9NamedExpr, Db9Row, Db9ScanKind, Db9SelectResponse, Db9Type, Db9TypeKind, Db9Value, }, kvrpcpb::{ApiVersion, IsolationLevel, LockInfo}, }; use matches::assert_matches; use more_asserts::{assert_ge, assert_lt}; use protobuf::Message; use raft::StateRole; use raftstore::{ RegionInfo, store::{ReadStats, WriteStats}, }; use resource_control::CpuThrottleManager; use tidb_query_common::util::convert_to_prefix_next; use tikv_kv::{MockEngine, MockEngineBuilder, destroy_tls_engine, set_tls_engine}; use tikv_util::config::ReadableSize; use tipb::{Executor, Expr}; use txn_types::{Key, LockType}; use super::*; use crate::{ config::{CoprReadPoolConfig, FairnessConfig, UnifiedReadPoolConfig}, coprocessor::readpool_impl::build_read_pool_for_test, read_pool::{ReadPool, build_yatp_read_pool}, storage::{ Store, TestEngineBuilder, kv::{BTreeEngine, FlowStatsReporter}, txn::tests::{must_commit, must_prewrite_put}, }, }; #[derive(Clone)] struct DummyReadPoolReporter; impl FlowStatsReporter for DummyReadPoolReporter { fn report_read_stats(&self, _read_stats: ReadStats) {} fn report_write_stats(&self, _write_stats: WriteStats) {} } fn build_endpoint_with_unified_yatp( keyspace_fairness_enabled: bool, end_point_max_concurrency: usize, ) -> Endpoint { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool_config = UnifiedReadPoolConfig { min_thread_count: 1, max_thread_count: 2, keyspace_fairness: FairnessConfig { enabled: keyspace_fairness_enabled, ..Default::default() }, ..Default::default() }; let read_pool = build_yatp_read_pool(&read_pool_config, DummyReadPoolReporter, engine, None); let server_config = Config { end_point_max_concurrency, ..Default::default() }; Endpoint::new( &server_config, read_pool.handle(), ConcurrencyManager::new(1.into()), None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ) } #[test] fn test_endpoint_uses_configured_concurrency_limit_without_fairness() { const MAX_CONCURRENCY: usize = 37; let endpoint = build_endpoint_with_unified_yatp(false, MAX_CONCURRENCY); let semaphore = endpoint .semaphore .expect("non-fair YATP should retain the concurrency limiter"); assert_eq!(semaphore.available_permits(), MAX_CONCURRENCY); } #[test] fn test_endpoint_bypasses_concurrency_limit_with_fairness() { let endpoint = build_endpoint_with_unified_yatp(true, 37); assert!(endpoint.semaphore.is_none()); } fn build_endpoint_with_cpu_throttle( cpu_throttle_config: resource_control::CpuThrottleConfig, global_cpu_us: u64, keyspace_cpu_us: std::collections::HashMap, read_pool_config: CoprReadPoolConfig, ) -> (Endpoint, BTreeEngine) { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test(&read_pool_config, engine.clone())); let cm = ConcurrencyManager::new(1.into()); let mut copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut resource_control_config = resource_control::Config::default(); resource_control_config.enabled = true; let mut controller = ResourceController::new(resource_control_config.clone()); let read_limiter = ReadLimiter::new(resource_control_config); let manager = Arc::new(CpuThrottleManager::new( cpu_throttle_config, global_cpu_us, keyspace_cpu_us, )); read_limiter.set_cpu_throttle_manager(manager); controller.set_read_limiter(read_limiter); copr.set_resource_controller(controller); (copr, engine) } fn make_req_ctx_with_priority( keyspace_id: u32, priority: kvrpcpb::CommandPri, deadline: Duration, ) -> ReqContext { let mut context = kvrpcpb::Context::default(); context.set_keyspace_id(keyspace_id); context.set_priority(priority); ReqContext::new( context, Vec::new(), deadline, None, None, txn_types::TimeStamp::max(), None, None, ) } fn make_batch_dag_request( keyspace_id: u32, priority: kvrpcpb::CommandPri, max_execution_duration_ms: u64, ) -> coppb::Request { let mut req = coppb::Request::default(); req.mut_context().set_isolation_level(IsolationLevel::Si); req.mut_context().set_keyspace_id(keyspace_id); req.mut_context().set_priority(priority); req.mut_context() .set_max_execution_duration_ms(max_execution_duration_ms); req.set_start_ts(100); req.set_tp(REQ_TYPE_DAG); let mut main_range = coppb::KeyRange::default(); main_range.set_start(tidb_query_datatype::codec::table::encode_row_key(1, 0)); main_range.set_end(tidb_query_datatype::codec::table::encode_row_key(1, 10)); req.mut_ranges().push(main_range); let mut scan = Executor::default(); scan.set_tp(ExecType::TypeTableScan); scan.mut_tbl_scan().set_table_id(1); let mut col = tipb::ColumnInfo::new(); col.set_column_id(1); scan.mut_tbl_scan().mut_columns().push(col); let mut dag = DagRequest::default(); dag.mut_executors().push(scan); req.set_data(dag.write_to_bytes().unwrap()); let mut batch_task = coppb::StoreBatchTask::new(); batch_task.set_task_id(42); let mut batch_range = coppb::KeyRange::default(); batch_range.set_start(tidb_query_datatype::codec::table::encode_row_key(1, 10)); batch_range.set_end(tidb_query_datatype::codec::table::encode_row_key(1, 20)); batch_task.mut_ranges().push(batch_range); req.mut_tasks().push(batch_task); req } fn run_tracked_unary_request( copr: Endpoint, engine: BTreeEngine, req_ctx: ReqContext, handler_builder: RequestHandlerBuilder<::Snap>, ) -> Result> { run_tracked_unary_request_with_fair_gate( copr, engine, req_ctx, handler_builder, false, None, ) } fn run_tracked_unary_request_with_fair_gate( copr: Endpoint, engine: BTreeEngine, req_ctx: ReqContext, handler_builder: RequestHandlerBuilder<::Snap>, fair_global_cpu_gate_active: bool, read_limiter_override: Option, ) -> Result> { run_tracked_unary_request_with_fair_gate_and_recorder( copr, engine, req_ctx, handler_builder, fair_global_cpu_gate_active, read_limiter_override, None, None, ) } fn run_tracked_unary_request_with_limiters( copr: Endpoint, engine: BTreeEngine, req_ctx: ReqContext, handler_builder: RequestHandlerBuilder<::Snap>, read_limiter_override: Option, keyspace_read_limiter_override: Option, ) -> Result> { run_tracked_unary_request_with_fair_gate_and_recorder( copr, engine, req_ctx, handler_builder, false, read_limiter_override, keyspace_read_limiter_override, None, ) } fn run_tracked_unary_request_with_fair_gate_and_recorder( copr: Endpoint, engine: BTreeEngine, req_ctx: ReqContext, handler_builder: RequestHandlerBuilder<::Snap>, fair_global_cpu_gate_active: bool, read_limiter_override: Option, keyspace_read_limiter_override: Option, poll_cpu_recorder: Option>, ) -> Result> { let keyspace_id = req_ctx.context.get_keyspace_id(); let tracker_token = GLOBAL_TRACKERS.insert(::tracker::Tracker::new(RequestInfo::new( &req_ctx.context, RequestType::CoprocessorDag, 0, ))); set_tls_tracker_token(tracker_token); let tracker = Box::new(Tracker::new(req_ctx, ReqTag::test, copr.slow_log_threshold)); set_tls_engine(engine); let read_limiter = read_limiter_override.or_else(|| copr.read_limiter.clone()); let keyspace_read_limiter = keyspace_read_limiter_override.or_else(|| { read_limiter .as_ref() .and_then(|limiter| limiter.get_limiter(keyspace_id)) }); let result = block_on(Endpoint::::handle_unary_request_impl( copr.semaphore.clone(), tracker, handler_builder, 0, None, copr.remote_ctx.clone(), tracker_token, copr.resource_publisher.clone(), poll_cpu_recorder, fair_global_cpu_gate_active, read_limiter, keyspace_read_limiter, false, false, )); unsafe { destroy_tls_engine::() }; clear_tls_tracker_token(); GLOBAL_TRACKERS.remove(tracker_token); result } fn wait_for_low_priority_bucket_available( copr: &Endpoint, expected_available: u64, ) { let manager = copr .read_limiter .as_ref() .and_then(|read_limiter| read_limiter.get_cpu_throttle_manager()) .expect("cpu throttle manager should exist"); let deadline = std::time::Instant::now() + Duration::from_secs(5); while std::time::Instant::now() < deadline { if manager.low_priority_dag_available() == expected_available { return; } thread::sleep(Duration::from_millis(5)); } panic!( "low-priority bucket did not reach expected availability {}, actual {}", expected_available, manager.low_priority_dag_available() ); } /// A unary `RequestHandler` that always produces a fixture. struct UnaryFixture { handle_duration_millis: u64, yieldable: bool, start_signal: Option>, runtime_budget_installed: Option>, result: Option>, } impl UnaryFixture { pub fn new(result: Result) -> UnaryFixture { UnaryFixture { handle_duration_millis: 0, yieldable: false, start_signal: None, runtime_budget_installed: None, result: Some(result), } } pub fn new_with_duration( result: Result, handle_duration_millis: u64, ) -> UnaryFixture { UnaryFixture { handle_duration_millis, yieldable: false, start_signal: None, runtime_budget_installed: None, result: Some(result), } } pub fn new_with_duration_and_start_signal( result: Result, handle_duration_millis: u64, start_signal: mpsc::Sender<()>, ) -> UnaryFixture { UnaryFixture { handle_duration_millis, yieldable: false, start_signal: Some(start_signal), runtime_budget_installed: None, result: Some(result), } } pub fn new_with_duration_yieldable( result: Result, handle_duration_millis: u64, ) -> UnaryFixture { UnaryFixture { handle_duration_millis, yieldable: true, start_signal: None, runtime_budget_installed: None, result: Some(result), } } pub fn new_with_runtime_budget_observer( result: Result, runtime_budget_installed: Arc, ) -> UnaryFixture { UnaryFixture { handle_duration_millis: 0, yieldable: false, start_signal: None, runtime_budget_installed: Some(runtime_budget_installed), result: Some(result), } } } #[async_trait] impl RequestHandler for UnaryFixture { async fn handle_request(&mut self) -> Result> { if let Some(start_signal) = self.start_signal.take() { let _ = start_signal.send(()); } if self.yieldable { // We split the task into small executions of 100 milliseconds. for _ in 0..self.handle_duration_millis / 100 { thread::sleep(Duration::from_millis(100)); yatp::task::future::reschedule().await; } thread::sleep(Duration::from_millis(self.handle_duration_millis % 100)); } else { thread::sleep(Duration::from_millis(self.handle_duration_millis)); } self.result.take().unwrap().map(|x| x.into()) } fn set_copr_runtime_budget( &mut self, budget: Option, ) -> bool { let Some(runtime_budget_installed) = &self.runtime_budget_installed else { return false; }; runtime_budget_installed.store(budget.is_some(), Ordering::SeqCst); true } } /// A streaming `RequestHandler` that always produces a fixture. struct StreamFixture { result_len: usize, result_iter: vec::IntoIter>, handle_durations_millis: vec::IntoIter, nth: usize, } impl StreamFixture { pub fn new(result: Vec>) -> StreamFixture { let len = result.len(); StreamFixture { result_len: len, result_iter: result.into_iter(), handle_durations_millis: vec![0; len].into_iter(), nth: 0, } } pub fn new_with_duration( result: Vec>, handle_durations_millis: Vec, ) -> StreamFixture { assert_eq!(result.len(), handle_durations_millis.len()); StreamFixture { result_len: result.len(), result_iter: result.into_iter(), handle_durations_millis: handle_durations_millis.into_iter(), nth: 0, } } } #[async_trait] impl RequestHandler for StreamFixture { async fn handle_streaming_request(&mut self) -> Result<(Option, bool)> { let is_finished = if self.result_len == 0 { true } else { self.nth >= (self.result_len - 1) }; let ret = match self.result_iter.next() { None => { assert!(is_finished); Ok((None, is_finished)) } Some(val) => { let handle_duration_ms = self.handle_durations_millis.next().unwrap(); thread::sleep(Duration::from_millis(handle_duration_ms)); match val { Ok(resp) => Ok((Some(resp), is_finished)), Err(e) => Err(e), } } }; self.nth += 1; ret } } /// A streaming `RequestHandler` that produces values according a closure. struct StreamFromClosure { result_generator: Box HandlerStreamStepResult + Send>, nth: usize, } impl StreamFromClosure { pub fn new(result_generator: F) -> StreamFromClosure where F: Fn(usize) -> HandlerStreamStepResult + Send + 'static, { StreamFromClosure { result_generator: Box::new(result_generator), nth: 0, } } } #[async_trait] impl RequestHandler for StreamFromClosure { async fn handle_streaming_request(&mut self) -> Result<(Option, bool)> { let result = (self.result_generator)(self.nth); self.nth += 1; result } } #[test] fn test_outdated_request() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); // a normal request let handler_builder = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let resp = block_on( copr.handle_unary_request(ParseCopRequestResult::default_for_test(handler_builder)), ) .unwrap(); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); // an outdated request let handler_builder = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let outdated_req_ctx = ReqContext::new( Default::default(), Vec::new(), Duration::from_secs(0), None, None, TimeStamp::max(), None, None, ); block_on(copr.handle_unary_request(ParseCopRequestResult { req_ctx: outdated_req_ctx, req_tag: ReqTag::test, remote_dispatch_precharge_bytes: 0, handler_builder, })) .unwrap_err(); } #[test] fn test_stack_guard() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let mut copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); copr.recursion_limit = 100; let req = { let mut expr = Expr::default(); for _ in 0..101 { let mut e = Expr::default(); e.mut_children().push(expr); expr = e; } let mut e = Executor::default(); e.mut_selection().mut_conditions().push(expr); let mut dag = DagRequest::default(); dag.mut_executors().push(e); let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DAG); req.set_data(dag.write_to_bytes().unwrap()); req }; let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!(!resp.get_other_error().is_empty()); } #[test] fn test_invalid_req_type() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut req = coppb::Request::default(); req.set_tp(9999); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!(!resp.get_other_error().is_empty()); } #[test] fn test_api_v3_request_is_rejected() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut req = coppb::Request::default(); req.mut_context().set_api_version(ApiVersion::V3); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!( resp.get_other_error() .contains("API V3 is not supported on the cloud-engine branch"), "unexpected response: {:?}", resp ); } #[test] fn test_invalid_req_body() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DAG); req.set_data(vec![1, 2, 3]); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!(!resp.get_other_error().is_empty()); } const DB9_TEST_REGION_ID: u64 = 23; const DB9_TEST_PEER_ID: u64 = 12; const DB9_TEST_STORE_ID: u64 = 100; fn build_db9_test_endpoint(engine: TestKvEngine) -> Endpoint { let (start_key, end_key) = db9_test_snapshot_range(&engine); let region_info_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![RegionInfo::new( metapb::Region { id: DB9_TEST_REGION_ID, start_key, end_key, peers: vec![metapb::Peer { id: DB9_TEST_PEER_ID, store_id: DB9_TEST_STORE_ID, ..Default::default() }] .into(), ..Default::default() }, StateRole::Leader, 0, )]); build_db9_test_endpoint_with_config( engine, Config::default(), ConcurrencyManager::new(1.into()), Some(region_info_accessor), ) } fn build_db9_test_endpoint_with_config( engine: E, config: Config, concurrency_manager: ConcurrencyManager, region_info_accessor: Option, ) -> Endpoint { let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); Endpoint::::new( &config, read_pool.handle(), concurrency_manager, region_info_accessor, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ) } fn make_db9_unary_request_with_ranges( dag: Db9DagRequest, ranges: Vec<(Vec, Vec)>, ) -> coppb::Request { let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DB9_DAG); req.set_start_ts(u64::MAX); req.mut_context().set_isolation_level(IsolationLevel::Si); req.mut_context().set_region_id(DB9_TEST_REGION_ID); req.mut_context().set_peer(metapb::Peer { id: DB9_TEST_PEER_ID, store_id: DB9_TEST_STORE_ID, ..Default::default() }); for (start, end) in ranges { let mut key_range = coppb::KeyRange::default(); key_range.set_start(start); key_range.set_end(end); req.mut_ranges().push(key_range); } req.set_data(dag.write_to_bytes().unwrap()); req } fn make_db9_unary_request(dag: Db9DagRequest, start: Vec, end: Vec) -> coppb::Request { make_db9_unary_request_with_ranges(dag, vec![(start, end)]) } fn make_remote_test_dag_request(engine: &TestKvEngine, scan_type: ExecType) -> coppb::Request { let mut scan = Executor::default(); scan.set_tp(scan_type); let mut col = tipb::ColumnInfo::new(); col.set_column_id(1); match scan_type { ExecType::TypeTableScan => { scan.mut_tbl_scan().set_table_id(1); scan.mut_tbl_scan().mut_columns().push(col); } ExecType::TypeIndexScan => { scan.mut_idx_scan().set_table_id(1); scan.mut_idx_scan().set_index_id(1); scan.mut_idx_scan().mut_columns().push(col); } _ => unreachable!("test only supports table and index scans"), } let mut dag = DagRequest::default(); dag.mut_executors().push(scan); let (start, end) = match scan_type { ExecType::TypeTableScan => ( tidb_query_datatype::codec::table::encode_row_key(1, 0), tidb_query_datatype::codec::table::encode_row_key(1, 1), ), ExecType::TypeIndexScan => { let start = tidb_query_datatype::codec::table::encode_index_seek_key(1, 1, &[]); let mut end = start.clone(); convert_to_prefix_next(&mut end); (start, end) } _ => unreachable!("test only supports table and index scans"), }; let snap = engine.engine.get_shard(1).unwrap().new_snap_access(); let keyspace_id = snap.get_keyspace_id(); let prefix = api_version::ApiV2::get_txn_keyspace_prefix(keyspace_id); let mut scoped_start = prefix.clone(); scoped_start.extend(start); let mut scoped_end = prefix; scoped_end.extend(end); let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DAG); req.set_start_ts(u64::MAX); req.mut_context().set_isolation_level(IsolationLevel::Si); req.mut_context().set_api_version(ApiVersion::V2); req.mut_context().set_keyspace_id(keyspace_id); req.mut_context().set_region_id(DB9_TEST_REGION_ID); req.mut_context().set_peer(metapb::Peer { id: DB9_TEST_PEER_ID, store_id: DB9_TEST_STORE_ID, ..Default::default() }); let mut key_range = coppb::KeyRange::default(); key_range.set_start(scoped_start); key_range.set_end(scoped_end); req.mut_ranges().push(key_range); req.set_data(dag.write_to_bytes().unwrap()); req } // Keep the full shard range and legacy request context so these cache tests // still exercise the remote dispatch policy on a cache miss. fn make_remote_cache_test_dag_request( engine: &TestKvEngine, cache_version: u64, ) -> coppb::Request { let mut scan = Executor::default(); scan.set_tp(ExecType::TypeTableScan); scan.mut_tbl_scan().set_table_id(1); let mut col = tipb::ColumnInfo::new(); col.set_column_id(1); scan.mut_tbl_scan().mut_columns().push(col); let mut dag = DagRequest::default(); dag.mut_executors().push(scan); let (start, end) = db9_test_snapshot_range(engine); let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DAG); req.set_start_ts(u64::MAX); req.set_is_cache_enabled(true); req.set_cache_if_match_version(cache_version); req.mut_context().set_isolation_level(IsolationLevel::Si); req.mut_context().set_region_id(DB9_TEST_REGION_ID); req.mut_context().set_peer(metapb::Peer { id: DB9_TEST_PEER_ID, store_id: DB9_TEST_STORE_ID, ..Default::default() }); let mut key_range = coppb::KeyRange::default(); key_range.set_start(start); key_range.set_end(end); req.mut_ranges().push(key_range); req.set_data(dag.write_to_bytes().unwrap()); req } fn make_db9_row_fetch_request( engine: &mut TestKvEngine, db_id: u64, table_id: u64, index_id: u64, ) -> (Db9DagRequest, Vec) { let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); put_db9_committed_row( engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &pk_bytes), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 120, 121, ); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); put_db9_committed_row(engine, index_key.clone(), pk_bytes, 130, 131); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], true)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("name_idx".to_owned()); dag.mut_scan().set_require_row_fetch(true); (dag, index_key) } fn run_remote_cop_result( req: coppb::Request, snap: ::Snap, ) -> Result { let tracker_token = GLOBAL_TRACKERS.insert(::tracker::Tracker::new(RequestInfo::new( req.get_context(), RequestType::CoprocessorDag, req.get_start_ts(), ))); set_tls_tracker_token(tracker_token); let result = block_on(parse_request_and_handle_remote_cop( req, None, Duration::from_secs(60), 1024, snap, )); clear_tls_tracker_token(); GLOBAL_TRACKERS.remove(tracker_token); result } fn run_remote_cop_request( req: coppb::Request, snap: ::Snap, ) -> Result> { run_remote_cop_result(req, snap).map(|result| result.response) } #[test] fn test_remote_cop_scan_stats_distinguish_table_and_index_requests() { for (scan_type, expected_kind) in [ (ExecType::TypeTableScan, RemoteCopRequestKind::Select), (ExecType::TypeIndexScan, RemoteCopRequestKind::Index), ] { let mut engine = TestKvEngine::new().unwrap(); let snap = block_on(crate::storage::Engine::async_snapshot( &mut engine, Default::default(), )) .unwrap(); let result = run_remote_cop_result(make_remote_test_dag_request(&engine, scan_type), snap) .unwrap(); assert_eq!(result.scan_stats.request_kind, expected_kind); } } #[test] fn test_cache_hit_skips_remote_coprocessor_dispatch() { let mut engine = TestKvEngine::new().unwrap(); let snap = block_on(crate::storage::Engine::async_snapshot( &mut engine, Default::default(), )) .unwrap(); let data_version = snap.ext().get_data_version().unwrap(); let mut copr = build_db9_test_endpoint(engine.clone()); copr.set_remote_url(Arc::new(VersionTrack::new(KvEngineConfig { remote_coprocessor_addr: "%%%invalid-dag-remote-uri%%%".to_owned(), remote_coprocessor_min_blocks_size: 0, remote_coprocessor_num_ranges: 1, remote_coprocessor_min_process_duration: tikv_util::config::ReadableDuration::millis(1), ..KvEngineConfig::default() }))); let resp = block_on(copr.parse_and_handle_unary_request( make_remote_cache_test_dag_request(&engine, data_version), None, )); assert!( resp.get_other_error().is_empty(), "cache hit should not reach invalid remote URL: {}", resp.get_other_error() ); assert!(resp.get_is_cache_hit()); assert_eq!(resp.get_cache_last_version(), data_version); assert!(resp.get_data().is_empty()); let remote_resp = block_on(copr.parse_and_handle_unary_request( make_remote_cache_test_dag_request(&engine, data_version.saturating_add(1)), None, )); assert!(!remote_resp.get_is_cache_hit()); assert!( remote_resp.get_other_error().contains("invalid"), "mismatched cache version should still exercise the remote path, got: {}", remote_resp.get_other_error() ); } #[test] fn test_stream_cache_match_does_not_use_copr_cache_hit() { let mut engine = TestKvEngine::new().unwrap(); let snap = block_on(crate::storage::Engine::async_snapshot( &mut engine, Default::default(), )) .unwrap(); let data_version = snap.ext().get_data_version().unwrap(); let mut copr = build_db9_test_endpoint(engine.clone()); copr.set_remote_url(Arc::new(VersionTrack::new(KvEngineConfig { remote_coprocessor_addr: "%%%invalid-dag-remote-uri%%%".to_owned(), remote_coprocessor_min_blocks_size: 0, remote_coprocessor_num_ranges: 1, remote_coprocessor_min_process_duration: tikv_util::config::ReadableDuration::millis(1), ..KvEngineConfig::default() }))); let responses = block_on_stream(copr.parse_and_handle_stream_request( make_remote_cache_test_dag_request(&engine, data_version), None, )) .collect::>(); assert_eq!(responses.len(), 1); assert!( !responses[0].get_is_cache_hit(), "streaming requests must not use the coprocessor cache-hit response" ); assert!( responses[0].get_other_error().contains("invalid"), "cache-valid streaming request should follow normal remote policy, got: {}", responses[0].get_other_error() ); } fn put_db9_committed_row( engine: &mut E, key: Vec, value: Vec, start_ts: u64, commit_ts: u64, ) { must_prewrite_put(engine, &key, &value, &key, start_ts); must_commit(engine, &key, start_ts, commit_ts); } fn db9_test_snapshot_range(engine: &TestKvEngine) -> (Vec, Vec) { // TestKvEngine snapshots expose shard outer bounds, so endpoint requests // must stay within that physical range even though DB9 keys are logical. let snap = engine.engine.get_shard(1).unwrap().new_snap_access(); ( snap.clone_start_key().to_vec(), snap.clone_end_key().to_vec(), ) } fn db9_test_scope_raw_key(engine: &TestKvEngine, raw_key: Vec) -> Vec { let snap = engine.engine.get_shard(1).unwrap().new_snap_access(); crate::coprocessor::db9::test_encode_snapshot_scoped_txn_raw_key( snap.clone_start_key().as_ref(), &raw_key, ) } fn db9_test_table_data_range(db_id: u64, table_id: u64) -> (Vec, Vec) { let start = crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[]); let mut end = start.clone(); convert_to_prefix_next(&mut end); (start, end) } fn db9_test_scoped_table_data_range( engine: &TestKvEngine, db_id: u64, table_id: u64, ) -> (Vec, Vec) { let (start, end) = db9_test_table_data_range(db_id, table_id); ( db9_test_scope_raw_key(engine, start), db9_test_scope_raw_key(engine, end), ) } fn db9_type(kind: Db9TypeKind) -> Db9Type { let mut out = Db9Type::default(); out.set_kind(kind); out } fn db9_column(name: &str, kind: Db9TypeKind) -> Db9ColumnInfo { let mut out = Db9ColumnInfo::default(); out.set_name(name.to_owned()); out.set_data_type(db9_type(kind)); out.set_nullable(false); out } fn db9_index(index_id: u64, index_name: &str, columns: &[&str], unique: bool) -> Db9IndexInfo { let mut out = Db9IndexInfo::default(); out.set_index_id(index_id); out.set_index_name(index_name.to_owned()); for column in columns { out.mut_columns().push((*column).to_owned()); } out.set_unique(unique); out } fn db9_column_expr(index: u32, name: &str, kind: Db9TypeKind) -> Db9Expr { let mut expr = Db9Expr::default(); expr.set_return_type(db9_type(kind)); let mut column = Db9ColumnRefExpr::default(); column.set_column_index(index); column.set_column_name(name.to_owned()); expr.set_column_ref(column); expr } fn db9_bool_const_expr(value: bool) -> Db9Expr { let mut expr = Db9Expr::default(); expr.set_return_type(db9_type(Db9TypeKind::Db9TypeKindBoolean)); let mut constant = Db9ConstantExpr::default(); let mut wire = Db9Value::default(); wire.set_bool_value(value); constant.set_value(wire); expr.set_constant(constant); expr } fn db9_text_const_expr(value: &str) -> Db9Expr { let mut expr = Db9Expr::default(); expr.set_return_type(db9_type(Db9TypeKind::Db9TypeKindText)); let mut constant = Db9ConstantExpr::default(); let mut wire = Db9Value::default(); wire.set_text_value(value.to_owned()); constant.set_value(wire); expr.set_constant(constant); expr } fn db9_int32_const_expr(value: i32) -> Db9Expr { let mut expr = Db9Expr::default(); expr.set_return_type(db9_type(Db9TypeKind::Db9TypeKindInt32)); let mut constant = Db9ConstantExpr::default(); let mut wire = Db9Value::default(); wire.set_int32_value(value); constant.set_value(wire); expr.set_constant(constant); expr } fn db9_func_expr(name: &str, return_kind: Db9TypeKind, args: Vec) -> Db9Expr { let mut expr = Db9Expr::default(); expr.set_return_type(db9_type(return_kind)); let mut func = Db9FuncCallExpr::default(); func.set_function_name(name.to_owned()); func.set_args(args.into()); expr.set_func_call(func); expr } fn db9_named_expr(name: &str, expr: Db9Expr) -> Db9NamedExpr { let mut out = Db9NamedExpr::default(); out.set_output_name(name.to_owned()); out.set_expr(expr); out } #[test] fn test_db9_unary_table_scan_response() { let db_id = 7; let table_id = 11; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), crate::coprocessor::db9::TestStoredValue::Boolean(true), ]), 10, 20, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 2]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), crate::coprocessor::db9::TestStoredValue::Text("bravo".to_owned()), crate::coprocessor::db9::TestStoredValue::Boolean(false), ]), 30, 40, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("c1", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("c2", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_columns() .push(db9_column("c3", Db9TypeKind::Db9TypeKindBoolean)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 2); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 2); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 2); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 2); assert_eq!(db9_resp.get_warnings().len(), 0); let first_row = &db9_resp.get_rows()[0]; assert_eq!(first_row.get_values().len(), 3); assert_eq!(first_row.get_values()[0].get_int64_value(), 1); assert_eq!(first_row.get_values()[1].get_text_value(), "alpha"); assert!(first_row.get_values()[2].get_bool_value()); let second_row = &db9_resp.get_rows()[1]; assert_eq!(second_row.get_values().len(), 3); assert_eq!(second_row.get_values()[0].get_int64_value(), 7); assert_eq!(second_row.get_values()[1].get_text_value(), "bravo"); assert!(!second_row.get_values()[2].get_bool_value()); } #[test] fn test_db9_unary_table_scan_api_v2_scoped_table_range_response() { let db_id = 7; let table_id = 11; let mut engine = TestKvEngine::new().unwrap(); let scoped_row_key = db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), ); put_db9_committed_row( &mut engine, scoped_row_key, crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 10, 20, ); let (range_start, range_end) = db9_test_scoped_table_data_range(&engine, db_id, table_id); let snap = engine.engine.get_shard(1).unwrap().new_snap_access(); let keyspace_id = snap.get_keyspace_id(); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("c1", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("c2", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut req = make_db9_unary_request(dag, range_start, range_end); req.mut_context().set_api_version(ApiVersion::V2); req.mut_context().set_keyspace_id(keyspace_id); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 1); assert_eq!( db9_resp.get_rows()[0].get_values()[1].get_text_value(), "alpha" ); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 1); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 1); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 1); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); } #[test] fn test_db9_table_scan_ignores_out_of_scope_keys_in_wide_snapshot_range() { let db_id = 8; let table_id = 12; let other_table_id = 13; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), ]), 10, 20, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, other_table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(999), ]), 30, 40, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected DB9 other_error: {}", resp.get_other_error() ); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 1); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 1); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 1); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 1); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); } #[test] fn test_db9_limit_is_applied() { let db_id = 9; let table_id = 21; let mut engine = TestKvEngine::new().unwrap(); for (index, (row_key, value)) in [ ([0_u8, 1_u8], 10_i64), ([0_u8, 2_u8], 20_i64), ([0_u8, 3_u8], 30_i64), ] .iter() .copied() .enumerate() { let start_ts = 50 + (index as u64 * 10); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &row_key), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(value), ]), start_ts, start_ts + 1, ); } let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); dag.mut_limit().set_limit(2); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 2); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 10); assert_eq!(db9_resp.get_rows()[1].get_values()[0].get_int64_value(), 20); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 2); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 2); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 2); } #[test] fn test_db9_filter_and_projection_are_applied() { let db_id = 12; let table_id = 34; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), crate::coprocessor::db9::TestStoredValue::Boolean(false), ]), 100, 101, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 2]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), crate::coprocessor::db9::TestStoredValue::Text("bravo".to_owned()), crate::coprocessor::db9::TestStoredValue::Boolean(true), ]), 110, 111, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_columns() .push(db9_column("active", Db9TypeKind::Db9TypeKindBoolean)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut selection = Db9Expr::default(); selection.set_return_type(db9_type(Db9TypeKind::Db9TypeKindBoolean)); let mut binary = Db9BinaryExpr::default(); binary.set_left(db9_column_expr( 2, "active", Db9TypeKind::Db9TypeKindBoolean, )); binary.set_op(Db9BinaryOp::Db9BinaryOpEq); binary.set_right(db9_bool_const_expr(true)); selection.set_binary(binary); dag.set_selection(selection); dag.mut_projections().push(db9_named_expr( "id", db9_column_expr(0, "id", Db9TypeKind::Db9TypeKindInt64), )); dag.mut_projections().push(db9_named_expr( "name", db9_column_expr(1, "name", Db9TypeKind::Db9TypeKindText), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 2); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 2); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); let row = &db9_resp.get_rows()[0]; assert_eq!(row.get_values().len(), 2); assert_eq!(row.get_values()[0].get_int64_value(), 2); assert_eq!(row.get_values()[1].get_text_value(), "bravo"); } #[test] fn test_db9_function_filter_and_projection_are_applied() { let db_id = 18; let table_id = 52; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("Alpha".to_owned()), crate::coprocessor::db9::TestStoredValue::Int64(-3), crate::coprocessor::db9::TestStoredValue::Null, ]), 180, 181, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 2]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), crate::coprocessor::db9::TestStoredValue::Text("Bravo".to_owned()), crate::coprocessor::db9::TestStoredValue::Int64(-7), crate::coprocessor::db9::TestStoredValue::Null, ]), 190, 191, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_columns() .push(db9_column("delta", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("note", Db9TypeKind::Db9TypeKindText)); let mut selection = Db9Expr::default(); selection.set_return_type(db9_type(Db9TypeKind::Db9TypeKindBoolean)); let mut binary = Db9BinaryExpr::default(); binary.set_left(db9_column_expr(1, "name", Db9TypeKind::Db9TypeKindText)); binary.set_op(Db9BinaryOp::Db9BinaryOpEq); binary.set_right(db9_text_const_expr("Bravo")); selection.set_binary(binary); dag.set_selection(selection); dag.mut_projections().push(db9_named_expr( "name_len", db9_func_expr( "length", Db9TypeKind::Db9TypeKindInt32, vec![db9_column_expr(1, "name", Db9TypeKind::Db9TypeKindText)], ), )); dag.mut_projections().push(db9_named_expr( "abs_delta", db9_func_expr( "abs", Db9TypeKind::Db9TypeKindInt64, vec![db9_column_expr(2, "delta", Db9TypeKind::Db9TypeKindInt64)], ), )); dag.mut_projections().push(db9_named_expr( "note_or_default", db9_func_expr( "coalesce", Db9TypeKind::Db9TypeKindText, vec![ db9_column_expr(3, "note", Db9TypeKind::Db9TypeKindText), db9_text_const_expr("fallback"), ], ), )); dag.mut_projections().push(db9_named_expr( "same_name_becomes_null", db9_func_expr( "nullif", Db9TypeKind::Db9TypeKindText, vec![ db9_column_expr(1, "name", Db9TypeKind::Db9TypeKindText), db9_text_const_expr("Bravo"), ], ), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 2); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 2); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); let row = &db9_resp.get_rows()[0]; assert_eq!(row.get_values().len(), 4); assert_eq!(row.get_values()[0].get_int32_value(), 5); assert_eq!(row.get_values()[1].get_int64_value(), 7); assert_eq!(row.get_values()[2].get_text_value(), "fallback"); assert!(row.get_values()[3].has_null_value()); } #[test] fn test_db9_index_scan_fetches_unique_entry_pk_from_value() { let db_id = 15; let table_id = 44; let index_id = 9; let mut engine = TestKvEngine::new().unwrap(); let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &pk_bytes), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 120, 121, ); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); put_db9_committed_row(&mut engine, index_key.clone(), pk_bytes, 130, 131); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], true)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("name_idx".to_owned()); dag.mut_scan().set_require_row_fetch(false); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 1); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 1); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 1); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 7); assert_eq!( db9_resp.get_rows()[0].get_values()[1].get_text_value(), "alpha" ); } #[test] fn test_db9_row_fetch_ignores_dag_region_accessor() { let db_id = 15; let table_id = 44; let index_id = 9; let mut engine = TestKvEngine::new().unwrap(); let (dag, index_key) = make_db9_row_fetch_request(&mut engine, db_id, table_id, index_id); let (start_key, end_key) = db9_test_snapshot_range(&engine); let region_info_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![RegionInfo::new( metapb::Region { id: DB9_TEST_REGION_ID, start_key, end_key, peers: vec![metapb::Peer { id: DB9_TEST_PEER_ID, store_id: DB9_TEST_STORE_ID, ..Default::default() }] .into(), ..Default::default() }, StateRole::Leader, 0, )]); let copr = build_db9_test_endpoint_with_config( engine.clone(), Config::default(), ConcurrencyManager::new(1.into()), Some(region_info_accessor), ); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert_eq!( resp.get_other_error(), "DB9 index scan row fetch requires extra region storage accessor" ); } #[test] fn test_db9_row_fetch_fails_without_extra_region_accessor() { let db_id = 15; let table_id = 44; let index_id = 9; let mut engine = TestKvEngine::new().unwrap(); let (dag, index_key) = make_db9_row_fetch_request(&mut engine, db_id, table_id, index_id); let copr = build_db9_test_endpoint_with_config( engine.clone(), Config::default(), ConcurrencyManager::new(1.into()), None, ); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert!( resp.get_other_error() .contains("DB9 index scan row fetch requires extra region storage accessor"), "unexpected other_error: {}", resp.get_other_error() ); } #[test] fn test_db9_index_only_scan_projects_index_key_and_pk_without_row_fetch() { let db_id = 18; let table_id = 47; let index_id = 12; let mut engine = TestKvEngine::new().unwrap(); let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], Some(&pk_bytes), ); put_db9_committed_row(&mut engine, index_key.clone(), vec![0x01], 130, 131); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_columns() .push(db9_column("payload", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("name_idx".to_owned()); dag.mut_scan().set_require_row_fetch(false); dag.mut_projections().push(db9_named_expr( "id", db9_column_expr(0, "id", Db9TypeKind::Db9TypeKindInt64), )); dag.mut_projections().push(db9_named_expr( "name", db9_column_expr(1, "name", Db9TypeKind::Db9TypeKindText), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 1); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 1); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 1); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 7); assert_eq!( db9_resp.get_rows()[0].get_values()[1].get_text_value(), "alpha" ); } #[test] fn test_db9_index_only_scan_rejects_uncovered_projection() { let db_id = 18; let table_id = 48; let index_id = 13; let mut engine = TestKvEngine::new().unwrap(); let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], Some(&pk_bytes), ); put_db9_committed_row(&mut engine, index_key.clone(), vec![0x01], 130, 131); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_columns() .push(db9_column("payload", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("name_idx".to_owned()); dag.mut_scan().set_require_row_fetch(false); dag.mut_projections().push(db9_named_expr( "payload", db9_column_expr(2, "payload", Db9TypeKind::Db9TypeKindText), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert!( resp.get_other_error() .contains("DB9 index-only scan references columns not covered"), "unexpected other_error: {}", resp.get_other_error() ); assert!(resp.get_other_error().contains("payload")); } #[test] fn test_db9_in_list_scan_fetches_key_suffix_pk_entries() { let db_id = 16; let table_id = 45; let index_id = 10; let mut engine = TestKvEngine::new().unwrap(); let alpha_pk = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), ]); let bravo_pk = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), ]); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &alpha_pk), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 140, 141, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &bravo_pk), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), crate::coprocessor::db9::TestStoredValue::Text("bravo".to_owned()), ]), 150, 151, ); let alpha_prefix = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); let bravo_prefix = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "bravo".to_owned(), )], None, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], Some(&alpha_pk), ), vec![0x01], 160, 161, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "bravo".to_owned(), )], Some(&bravo_pk), ), vec![0x01], 170, 171, ); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindInList); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("name_idx".to_owned()); dag.mut_scan().set_require_row_fetch(false); dag.mut_projections().push(db9_named_expr( "id", db9_column_expr(0, "id", Db9TypeKind::Db9TypeKindInt64), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request_with_ranges( dag, vec![ ( db9_test_scope_raw_key(&engine, alpha_prefix.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&alpha_prefix), ), ), ( db9_test_scope_raw_key(&engine, bravo_prefix.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&bravo_prefix), ), ), ], ), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 2); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 2); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 2); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 2); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 1); assert_eq!(db9_resp.get_rows()[1].get_values()[0].get_int64_value(), 2); } #[test] fn test_db9_desc_index_scan_preserves_global_order_across_ranges() { let db_id = 19; let table_id = 49; let index_id = 14; let mut engine = TestKvEngine::new().unwrap(); let first_pk = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), ]); let second_pk = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), ]); let first_prefix = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Int64(100)], None, ); let second_prefix = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Int64(200)], None, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Int64(100)], Some(&first_pk), ), vec![0x01], 180, 181, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Int64(200)], Some(&second_pk), ), vec![0x01], 190, 191, ); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("created_at", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table().mut_indexes().push(db9_index( index_id, "created_at_idx", &["created_at"], false, )); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndexRange); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("created_at_idx".to_owned()); dag.mut_scan().set_require_row_fetch(false); dag.mut_scan().set_desc(true); dag.mut_projections().push(db9_named_expr( "id", db9_column_expr(0, "id", Db9TypeKind::Db9TypeKindInt64), )); dag.mut_projections().push(db9_named_expr( "created_at", db9_column_expr(1, "created_at", Db9TypeKind::Db9TypeKindInt64), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request_with_ranges( dag, vec![ ( db9_test_scope_raw_key(&engine, first_prefix.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&first_prefix), ), ), ( db9_test_scope_raw_key(&engine, second_prefix.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&second_prefix), ), ), ], ), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!(!resp.has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 2); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 2); assert_eq!( db9_resp.get_rows()[0].get_values()[1].get_int64_value(), 200 ); assert_eq!(db9_resp.get_rows()[1].get_values()[0].get_int64_value(), 1); assert_eq!( db9_resp.get_rows()[1].get_values()[1].get_int64_value(), 100 ); } #[test] fn test_db9_index_scan_ignores_out_of_scope_keys_in_wide_snapshot_range() { let db_id = 17; let table_id = 46; let index_id = 11; let other_index_id = 12; let mut engine = TestKvEngine::new().unwrap(); let target_pk = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); let other_pk = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(9), ]); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &target_pk), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 120, 121, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &other_pk), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(9), crate::coprocessor::db9::TestStoredValue::Text("bravo".to_owned()), ]), 130, 131, ); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); put_db9_committed_row(&mut engine, index_key, target_pk, 140, 141); let other_index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, other_index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "bravo".to_owned(), )], None, ); put_db9_committed_row(&mut engine, other_index_key, other_pk, 150, 151); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], true)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_index_name("name_idx".to_owned()); dag.mut_scan().set_require_row_fetch(false); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected DB9 other_error: {}", resp.get_other_error() ); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 7); assert_eq!( db9_resp.get_rows()[0].get_values()[1].get_text_value(), "alpha" ); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 1); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 1); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 1); assert_eq!(db9_resp.get_stats().get_returned_rows(), 1); } #[test] fn test_db9_local_request_does_not_bypass_all_locks() { let db_id = 18; let table_id = 47; let raw_key = crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]); let engine = TestKvEngine::new().unwrap(); let key = db9_test_scope_raw_key(&engine, raw_key.clone()); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut req = make_db9_unary_request( dag, key.clone(), db9_test_scope_raw_key(&engine, crate::coprocessor::db9::test_prefix_end(&raw_key)), ); req.mut_context().set_isolation_level(IsolationLevel::Si); req.mut_context().mut_resolved_locks().push(42); let parsed = copr .parse_request_and_check_memory_locks(req, None, false, u64::MAX) .unwrap(); assert_eq!( parsed.req_ctx.bypass_locks, txn_types::TsSet::from_u64s(vec![42]) ); assert_ne!(parsed.req_ctx.bypass_locks, txn_types::TsSet::All); } #[test] fn test_db9_index_scan_rejects_non_plain_local_si_context() { let db_id = 19; let table_id = 48; let index_id = 3; let mut engine = TestKvEngine::new().unwrap(); let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); put_db9_committed_row(&mut engine, index_key.clone(), pk_bytes, 130, 131); let copr = build_db9_test_endpoint(engine.clone()); let range_start = db9_test_scope_raw_key(&engine, index_key.clone()); let range_end = db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ); let build_req = || { let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_require_row_fetch(true); make_db9_unary_request(dag, range_start.clone(), range_end.clone()) }; let mut stale_req = build_req(); stale_req.mut_context().set_stale_read(true); let stale_resp = block_on(copr.parse_and_handle_unary_request(stale_req, None)); assert!( stale_resp .get_other_error() .contains("DB9 index scan row fetch only supports plain local SI requests"), "unexpected stale_read error: {}", stale_resp.get_other_error() ); let mut replica_req = build_req(); replica_req.mut_context().set_replica_read(true); let replica_resp = block_on(copr.parse_and_handle_unary_request(replica_req, None)); assert!( replica_resp .get_other_error() .contains("DB9 index scan row fetch only supports plain local SI requests"), "unexpected replica_read error: {}", replica_resp.get_other_error() ); let mut rc_req = build_req(); rc_req.mut_context().set_isolation_level(IsolationLevel::Rc); let rc_resp = block_on(copr.parse_and_handle_unary_request(rc_req, None)); assert!( rc_resp .get_other_error() .contains("DB9 index scan row fetch only supports plain local SI requests"), "unexpected non-SI error: {}", rc_resp.get_other_error() ); } #[test] fn test_db9_invalid_req_body() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DB9_DAG); req.set_data(vec![1, 2, 3]); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!(!resp.get_other_error().is_empty()); } #[test] fn test_db9_rejects_missing_codec_version() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut dag = Db9DagRequest::default(); dag.set_db_id(1); dag.mut_table().set_table_id(1); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut req = coppb::Request::default(); req.set_tp(REQ_TYPE_DB9_DAG); req.set_start_ts(u64::MAX); let mut range = coppb::KeyRange::default(); range.set_start(b"a".to_vec()); range.set_end(b"z".to_vec()); req.mut_ranges().push(range); req.set_data(dag.write_to_bytes().unwrap()); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!( resp.get_other_error() .contains("unsupported DB9 codec_version 0, expected 3"), "unexpected other_error: {}", resp.get_other_error() ); } #[test] fn test_db9_streaming_request_returns_rows() { let db_id = 31; let table_id = 41; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 10, 20, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 2]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), crate::coprocessor::db9::TestStoredValue::Text("bravo".to_owned()), ]), 30, 40, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let req = make_db9_unary_request(dag, range_start, range_end); let responses = block_on_stream(copr.parse_and_handle_stream_request(req, None)).collect::>(); assert_eq!(responses.len(), 1); assert!( responses[0].get_other_error().is_empty(), "unexpected other_error: {}", responses[0].get_other_error() ); let mut db9_resp = Db9SelectResponse::default(); db9_resp .merge_from_bytes(responses[0].get_data()) .expect("decode streamed DB9 response"); assert_eq!(db9_resp.get_rows().len(), 2); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 1); assert_eq!(db9_resp.get_rows()[1].get_values()[0].get_int64_value(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 2); } #[test] fn test_db9_streaming_api_v2_scoped_table_range_returns_rows() { let db_id = 31; let table_id = 41; let mut engine = TestKvEngine::new().unwrap(); for (pk, id, name) in [(1_u8, 1_i64, "alpha"), (2_u8, 2_i64, "bravo")] { let scoped_row_key = db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, pk]), ); put_db9_committed_row( &mut engine, scoped_row_key, crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(id), crate::coprocessor::db9::TestStoredValue::Text(name.to_owned()), ]), (10 + id) as u64, (20 + id) as u64, ); } let (range_start, range_end) = db9_test_scoped_table_data_range(&engine, db_id, table_id); let snap = engine.engine.get_shard(1).unwrap().new_snap_access(); let keyspace_id = snap.get_keyspace_id(); let copr = build_db9_test_endpoint(engine); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut req = make_db9_unary_request(dag, range_start, range_end); req.mut_context().set_api_version(ApiVersion::V2); req.mut_context().set_keyspace_id(keyspace_id); let responses = block_on_stream(copr.parse_and_handle_stream_request(req, None)).collect::>(); assert_eq!(responses.len(), 1); assert!( responses[0].get_other_error().is_empty(), "unexpected other_error: {}", responses[0].get_other_error() ); assert!(!responses[0].has_region_error()); let mut db9_resp = Db9SelectResponse::default(); db9_resp .merge_from_bytes(responses[0].get_data()) .expect("decode streamed DB9 response"); assert_eq!(db9_resp.get_rows().len(), 2); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 1); assert_eq!(db9_resp.get_rows()[1].get_values()[0].get_int64_value(), 2); assert_eq!(db9_resp.get_stats().get_scanned_keys(), 2); assert_eq!(db9_resp.get_stats().get_scanned_rows(), 2); assert_eq!(db9_resp.get_stats().get_decoded_rows(), 2); assert_eq!(db9_resp.get_stats().get_returned_rows(), 2); } #[test] fn test_db9_request_uses_remote_cop_when_configured() { let db_id = 21; let table_id = 51; let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]); let mut engine = TestKvEngine::new().unwrap(); let row_key = crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &pk_bytes); put_db9_committed_row( &mut engine, row_key, crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), ]), 120, 121, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let mut copr = build_db9_test_endpoint(engine.clone()); copr.set_remote_url(Arc::new(VersionTrack::new(KvEngineConfig { remote_worker_addr: String::new(), remote_coprocessor_addr: "%%%invalid-db9-remote-uri%%%".to_owned(), remote_coprocessor_min_blocks_size: 0, remote_coprocessor_num_ranges: 1, remote_coprocessor_min_process_duration: tikv_util::config::ReadableDuration::millis(1), ..KvEngineConfig::default() }))); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().contains("invalid"), "unexpected DB9 remote error: {}", resp.get_other_error() ); } #[test] fn test_db9_remote_worker_rejects_index_scan_requests() { let db_id = 22; let table_id = 52; let index_id = 9; let mut engine = TestKvEngine::new().unwrap(); let (range_start, range_end) = db9_test_snapshot_range(&engine); let snap = block_on(crate::storage::Engine::async_snapshot( &mut engine, Default::default(), )) .unwrap(); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_require_row_fetch(false); let err = run_remote_cop_request(make_db9_unary_request(dag, range_start, range_end), snap) .unwrap_err(); assert!( err.to_string() .contains("remote DB9 cop does not support index scan requests"), "unexpected remote DB9 index-scan error: {}", err, ); } #[test] fn test_db9_remote_worker_rejects_limit_requests() { let db_id = 23; let table_id = 53; let mut engine = TestKvEngine::new().unwrap(); let (range_start, range_end) = db9_test_snapshot_range(&engine); let snap = block_on(crate::storage::Engine::async_snapshot( &mut engine, Default::default(), )) .unwrap(); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); dag.mut_limit().set_limit(1); let err = run_remote_cop_request(make_db9_unary_request(dag, range_start, range_end), snap) .unwrap_err(); assert!( err.to_string() .contains("remote DB9 cop does not support pushed LIMIT requests"), "unexpected remote DB9 limit error: {}", err, ); } #[test] fn test_db9_request_stays_local_when_remote_load_aware_detects_idle_cpu() { let db_id = 22; let table_id = 52; let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(8), ]); let mut engine = TestKvEngine::new().unwrap(); let row_key = crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &pk_bytes); put_db9_committed_row( &mut engine, row_key, crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(8), ]), 220, 221, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let mut copr = build_db9_test_endpoint(engine.clone()); let mut resource_control_config = resource_control::Config::default(); resource_control_config.enabled = true; let mut controller = ResourceController::new(resource_control_config.clone()); controller.set_read_limiter(ReadLimiter::new(resource_control_config)); let kvengine_config = Arc::new(VersionTrack::new(KvEngineConfig { remote_worker_addr: String::new(), remote_coprocessor_addr: "%%%invalid-db9-remote-uri%%%".to_owned(), remote_coprocessor_min_blocks_size: 0, remote_coprocessor_num_ranges: 1, remote_coprocessor_min_process_duration: tikv_util::config::ReadableDuration::millis(1), remote_coprocessor_load_aware: true, remote_coprocessor_load_aware_idle_threshold: 0.3, remote_coprocessor_load_aware_busy_threshold: 0.5, remote_coprocessor_load_aware_force_blocks_size: 256 * 1024 * 1024, remote_coprocessor_load_aware_force_num_ranges: 2, ..KvEngineConfig::default() })); let total_cpu_monitor = controller.total_cpu_monitor(); copr.configure_remote_coprocessor(kvengine_config, controller); total_cpu_monitor.set_ratio(0.20); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected DB9 local error: {}", resp.get_other_error() ); assert!( !resp.get_data().is_empty(), "expected DB9 request to execute locally" ); } #[test] fn test_db9_success_metrics_are_recorded() { let db_id = 31; let table_id = 41; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(11), ]), 10, 11, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 2]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(22), ]), 20, 21, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); } #[test] fn test_db9_request_cache_is_ignored_even_with_matching_version() { let db_id = 32; let table_id = 42; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(11), ]), 10, 11, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut req = make_db9_unary_request(dag, range_start, range_end); req.set_is_cache_enabled(true); req.set_cache_if_match_version(42); let parsed = copr .parse_request_and_check_memory_locks(req.clone(), None, false, u64::MAX) .unwrap(); assert!( parsed.req_ctx.cache_match_version.is_none(), "DB9 request should ignore generic request-cache metadata" ); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); assert!( !resp.get_is_cache_hit(), "DB9 request must not hit generic cache path" ); let mut db9_resp = Db9SelectResponse::default(); db9_resp.merge_from_bytes(resp.get_data()).unwrap(); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 11); } #[test] fn test_db9_ru_v2_metrics_are_recorded() { let db_id = 33; let table_id = 43; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(11), ]), 10, 11, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("value", Db9TypeKind::Db9TypeKindInt64)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); dag.set_selection(db9_bool_const_expr(true)); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); let ru_v2 = resp.get_exec_details_v2().get_ru_v2(); assert_eq!( ru_v2 .get_executor_inputs() .get_tikv_coprocessor_executor_work_total_batch_table_scan(), 1 ); assert_eq!( ru_v2 .get_executor_inputs() .get_tikv_coprocessor_executor_work_total_batch_selection(), 1 ); assert_eq!(ru_v2.get_coprocessor_executor_iterations(), 2); } #[test] fn test_db9_index_only_scan_ru_v2_metrics_include_selection() { let db_id = 34; let table_id = 44; let index_id = 8; let mut engine = TestKvEngine::new().unwrap(); let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(9), ]); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &pk_bytes), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(9), crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 10, 11, ); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); put_db9_committed_row(&mut engine, index_key.clone(), pk_bytes, 130, 131); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table().mut_pk_indices().push(0); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_require_row_fetch(false); dag.set_selection(db9_bool_const_expr(true)); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert!( resp.get_other_error().is_empty(), "unexpected other_error: {}", resp.get_other_error() ); let ru_v2 = resp.get_exec_details_v2().get_ru_v2(); assert_eq!( ru_v2 .get_executor_inputs() .get_tikv_coprocessor_executor_work_total_batch_index_scan(), 0 ); assert_eq!( ru_v2 .get_executor_inputs() .get_tikv_coprocessor_executor_work_total_batch_table_scan(), 1 ); assert_eq!( ru_v2 .get_executor_inputs() .get_tikv_coprocessor_executor_work_total_batch_selection(), 1 ); assert_eq!(ru_v2.get_coprocessor_executor_iterations(), 2); } #[test] fn test_db9_index_scan_missing_row_fails_closed() { let db_id = 35; let table_id = 45; let index_id = 9; let mut engine = TestKvEngine::new().unwrap(); let pk_bytes = crate::coprocessor::db9::test_encode_pk_bytes(vec![ crate::coprocessor::db9::TestStoredValue::Int64(11), ]); let index_key = crate::coprocessor::db9::test_encode_index_key( db_id, table_id, index_id, vec![crate::coprocessor::db9::TestStoredValue::Text( "alpha".to_owned(), )], None, ); put_db9_committed_row(&mut engine, index_key.clone(), pk_bytes, 130, 131); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_table() .mut_indexes() .push(db9_index(index_id, "name_idx", &["name"], false)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex); dag.mut_scan().set_index_id(index_id); dag.mut_scan().set_require_row_fetch(true); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request( dag, db9_test_scope_raw_key(&engine, index_key.clone()), db9_test_scope_raw_key( &engine, crate::coprocessor::db9::test_prefix_end(&index_key), ), ), None, )); assert!( resp.get_other_error() .contains("DB9 index scan row fetch requires extra region storage accessor"), "unexpected other_error: {}", resp.get_other_error() ); } #[test] fn test_db9_projection_eval_errors_are_recorded() { let db_id = 32; let table_id = 42; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Text("alpha".to_owned()), ]), 10, 11, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let copr = build_db9_test_endpoint(engine.clone()); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("name", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); dag.mut_projections().push(db9_named_expr( "unsupported", db9_func_expr( "substr", Db9TypeKind::Db9TypeKindText, vec![ db9_column_expr(0, "name", Db9TypeKind::Db9TypeKindText), db9_int32_const_expr(1), db9_int32_const_expr(-1), ], ), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error() .contains("negative substring length not allowed"), "unexpected other_error: {}", resp.get_other_error() ); } #[test] fn test_db9_unary_response_respects_max_resp_size() { let db_id = 33; let table_id = 43; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Text("x".repeat(256)), ]), 10, 11, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let mut config = Config::default(); config.cop_max_resp_size = tikv_util::config::ReadableSize(64); let copr = build_db9_test_endpoint_with_config( engine.clone(), config, ConcurrencyManager::new(1.into()), None, ); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("payload", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error() .contains("exceeds coprocessor max_resp_size"), "unexpected DB9 other_error: {}", resp.get_other_error() ); } #[test] fn test_db9_streaming_splits_chunks_under_max_resp_size() { let db_id = 35; let table_id = 45; let payloads = (1..=3) .map(|id| format!("payload-{id:02}-{}", "x".repeat(48))) .collect::>(); let wire_row = |id: i64, payload: &str| { let mut row = Db9Row::default(); let mut id_value = Db9Value::default(); id_value.set_int64_value(id); row.mut_values().push(id_value); let mut payload_value = Db9Value::default(); payload_value.set_text_value(payload.to_owned()); row.mut_values().push(payload_value); row }; let response_len = |rows: Vec, returned_rows: u64| { let mut response = Db9SelectResponse::default(); for row in rows { response.mut_rows().push(row); } response.mut_stats().set_scanned_keys(3); response.mut_stats().set_scanned_rows(3); response.mut_stats().set_decoded_rows(3); response.mut_stats().set_returned_rows(returned_rows); response.write_to_bytes().unwrap().len() }; let one_row_frame_size = response_len(vec![wire_row(1, &payloads[0])], 1); let two_row_frame_size = response_len( vec![wire_row(1, &payloads[0]), wire_row(2, &payloads[1])], 2, ); let max_resp_size = one_row_frame_size + 1; assert!( max_resp_size < two_row_frame_size, "test setup should allow one DB9 row per stream frame" ); let mut engine = TestKvEngine::new().unwrap(); for id in 1_i64..=3 { put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key( db_id, table_id, &[0, id as u8], ), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(id), crate::coprocessor::db9::TestStoredValue::Text( payloads[(id - 1) as usize].clone(), ), ]), (10 + id) as u64, (20 + id) as u64, ); } let (range_start, range_end) = db9_test_snapshot_range(&engine); let mut config = Config::default(); config.cop_max_resp_size = tikv_util::config::ReadableSize(max_resp_size as u64); let copr = build_db9_test_endpoint_with_config( engine, config, ConcurrencyManager::new(1.into()), None, ); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("payload", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let responses = block_on_stream(copr.parse_and_handle_stream_request( make_db9_unary_request(dag, range_start, range_end), None, )) .collect::>(); assert!( responses.len() > 1, "expected multiple DB9 stream chunks, got {}", responses.len() ); let mut ids = Vec::new(); for response in responses { assert!( response.get_other_error().is_empty(), "unexpected DB9 other_error: {}", response.get_other_error() ); assert!( response.get_data().len() <= max_resp_size, "stream frame exceeded configured max_resp_size: {}", response.get_data().len() ); let mut db9_resp = Db9SelectResponse::default(); db9_resp .merge_from_bytes(response.get_data()) .expect("decode streamed DB9 response chunk"); ids.extend( db9_resp .get_rows() .iter() .map(|row| row.get_values()[0].get_int64_value()), ); } assert_eq!(ids, vec![1, 2, 3]); } #[test] fn test_db9_streaming_reports_error_after_emitted_chunk() { let db_id = 36; let table_id = 46; let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(1), crate::coprocessor::db9::TestStoredValue::Text("ok".to_owned()), ]), 10, 20, ); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 2]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(2), crate::coprocessor::db9::TestStoredValue::Text("x".repeat(60)), ]), 30, 40, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let mut config = Config::default(); config.cop_max_resp_size = tikv_util::config::ReadableSize(64); let copr = build_db9_test_endpoint_with_config( engine, config, ConcurrencyManager::new(1.into()), None, ); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("payload", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let responses = block_on_stream(copr.parse_and_handle_stream_request( make_db9_unary_request(dag, range_start, range_end), None, )) .collect::>(); assert_eq!(responses.len(), 2); assert!( responses[0].get_other_error().is_empty(), "unexpected first DB9 other_error: {}", responses[0].get_other_error() ); let mut db9_resp = Db9SelectResponse::default(); db9_resp .merge_from_bytes(responses[0].get_data()) .expect("decode first streamed DB9 response chunk"); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 1); assert!( responses[1] .get_other_error() .contains("exceeds coprocessor max_resp_size"), "unexpected DB9 other_error: {}", responses[1].get_other_error() ); } #[test] fn test_db9_wide_filtered_row_survives_small_max_resp_size() { let db_id = 34; let table_id = 44; let payload = "x".repeat(256); let mut engine = TestKvEngine::new().unwrap(); put_db9_committed_row( &mut engine, crate::coprocessor::db9::test_encode_table_data_key(db_id, table_id, &[0, 1]), crate::coprocessor::db9::test_encode_row(vec![ crate::coprocessor::db9::TestStoredValue::Int64(7), crate::coprocessor::db9::TestStoredValue::Text(payload.clone()), ]), 10, 11, ); let (range_start, range_end) = db9_test_snapshot_range(&engine); let mut config = Config::default(); config.cop_max_resp_size = tikv_util::config::ReadableSize(64); let copr = build_db9_test_endpoint_with_config( engine.clone(), config, ConcurrencyManager::new(1.into()), None, ); let mut dag = Db9DagRequest::default(); dag.set_codec_version(crate::coprocessor::db9::DB9_CODEC_VERSION); dag.set_db_id(db_id); dag.mut_table().set_table_id(table_id); dag.mut_table() .mut_columns() .push(db9_column("id", Db9TypeKind::Db9TypeKindInt64)); dag.mut_table() .mut_columns() .push(db9_column("payload", Db9TypeKind::Db9TypeKindText)); dag.mut_scan().set_kind(Db9ScanKind::Db9ScanKindTable); let mut selection = Db9Expr::default(); selection.set_return_type(db9_type(Db9TypeKind::Db9TypeKindBoolean)); let mut binary = Db9BinaryExpr::default(); binary.set_left(db9_func_expr( "reverse", Db9TypeKind::Db9TypeKindText, vec![db9_column_expr(1, "payload", Db9TypeKind::Db9TypeKindText)], )); binary.set_op(Db9BinaryOp::Db9BinaryOpEq); binary.set_right(db9_text_const_expr(&payload)); selection.set_binary(binary); dag.set_selection(selection); dag.mut_projections().push(db9_named_expr( "id", db9_column_expr(0, "id", Db9TypeKind::Db9TypeKindInt64), )); let resp = block_on(copr.parse_and_handle_unary_request( make_db9_unary_request(dag, range_start, range_end), None, )); assert!( resp.get_other_error().is_empty(), "unexpected DB9 other_error: {}", resp.get_other_error() ); let mut db9_resp = Db9SelectResponse::default(); db9_resp .merge_from_bytes(resp.get_data()) .expect("decode DB9 response"); assert_eq!(db9_resp.get_rows().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values().len(), 1); assert_eq!(db9_resp.get_rows()[0].get_values()[0].get_int64_value(), 7); } #[test] fn test_remote_cop_request_type() { let mut req = coppb::Request::default(); for (tp, expected) in [ (REQ_TYPE_DAG, RemoteCopRequestType::Dag), (REQ_TYPE_DB9_DAG, RemoteCopRequestType::Db9), (REQ_TYPE_ANALYZE, RemoteCopRequestType::Analyze), (REQ_TYPE_CHECKSUM, RemoteCopRequestType::Checksum), (0, RemoteCopRequestType::Unknown), ] { req.set_tp(tp); assert_eq!(remote_cop_request_type(&req), expected); } } #[test] fn test_db9_scan_kind_uses_dedicated_req_tags() { assert_eq!( db9_req_tag_from_scan_kind(kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindTable), ReqTag::db9_select ); assert_eq!( db9_req_tag_from_scan_kind(kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindIndex), ReqTag::db9_index ); assert_eq!( db9_req_tag_from_scan_kind( kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindIndexRange ), ReqTag::db9_index ); assert_eq!( db9_req_tag_from_scan_kind( kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindIndexBoundedRange ), ReqTag::db9_index ); assert_eq!( db9_req_tag_from_scan_kind(kvproto::db9_coprocessor::Db9ScanKind::Db9ScanKindInList), ReqTag::db9_index ); } #[test] fn test_full() { use std::sync::Mutex; use tikv_util::yatp_pool::{DefaultTicker, YatpPoolBuilder}; use crate::storage::kv::{destroy_tls_engine, set_tls_engine}; let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from( CoprReadPoolConfig { normal_concurrency: 1, max_tasks_per_worker_normal: 2, ..CoprReadPoolConfig::default_for_test() } .to_yatp_pool_configs() .into_iter() .map(|config| { let engine = Arc::new(Mutex::new(engine.clone())); YatpPoolBuilder::new(DefaultTicker) .config(config) .name_prefix("coprocessor_endpoint_test_full") .after_start(move || set_tls_engine(engine.lock().unwrap().clone())) // Safety: we call `set_` and `destroy_` with the same engine type. .before_stop(|| unsafe { destroy_tls_engine::() }) .build_future_pool() }) .collect::>(), ); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let (tx, rx) = mpsc::channel(); // first 2 requests are processed as normal and laters are returned as errors for i in 0..5 { let mut response = coppb::Response::default(); response.set_data(vec![1, 2, i]); let mut context = kvrpcpb::Context::default(); context.set_priority(kvrpcpb::CommandPri::Normal); let handler_builder = Box::new(|_, _: &_| { Ok(UnaryFixture::new_with_duration(Ok(response), 1000).into_boxed()) }); let future = copr.handle_unary_request(ParseCopRequestResult::default_for_test(handler_builder)); let tx = tx.clone(); thread::spawn(move || { tx.send(block_on(future)).unwrap(); }); thread::sleep(Duration::from_millis(100)); } // verify for _ in 2..5 { rx.recv().unwrap().unwrap_err(); } for i in 0..2 { let resp = rx.recv().unwrap().unwrap(); assert_eq!(resp.get_data(), [1, 2, i]); assert!(!resp.has_region_error()); } } #[test] fn test_error_unary_response() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let handler_builder = Box::new(|_, _: &_| Ok(UnaryFixture::new(Err(box_err!("foo"))).into_boxed())); let resp = block_on( copr.handle_unary_request(ParseCopRequestResult::default_for_test(handler_builder)), ) .unwrap(); assert_eq!(resp.get_data().len(), 0); assert!(!resp.get_other_error().is_empty()); } #[test] fn test_error_streaming_response() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); // Fail immediately let handler_builder = Box::new(|_, _: &_| Ok(StreamFixture::new(vec![Err(box_err!("foo"))]).into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 1); assert_eq!(resp_vec[0].get_data().len(), 0); assert!(!resp_vec[0].get_other_error().is_empty()); // Fail after some success responses let mut responses = Vec::new(); for i in 0..5 { let mut resp = coppb::Response::default(); resp.set_data(vec![1, 2, i]); responses.push(Ok(resp)); } responses.push(Err(box_err!("foo"))); let handler_builder = Box::new(|_, _: &_| Ok(StreamFixture::new(responses).into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 6); for (i, resp) in resp_vec.iter().enumerate().take(5) { assert_eq!(resp.get_data(), [1, 2, i as u8]); } assert_eq!(resp_vec[5].get_data().len(), 0); assert!(!resp_vec[5].get_other_error().is_empty()); } #[test] fn test_empty_streaming_response() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let handler_builder = Box::new(|_, _: &_| Ok(StreamFixture::new(vec![]).into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 0); } // TODO: Test panic? #[test] fn test_special_streaming_handlers() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); // handler returns `finished == true` should not be called again. let counter = Arc::new(atomic::AtomicIsize::new(0)); let counter_clone = Arc::clone(&counter); let handler = StreamFromClosure::new(move |nth| match nth { 0 => { let mut resp = coppb::Response::default(); resp.set_data(vec![1, 2, 7]); Ok((Some(resp), true)) } _ => { // we cannot use `unreachable!()` here because CpuPool catches panic. counter_clone.store(1, atomic::Ordering::SeqCst); Err(box_err!("unreachable")) } }); let handler_builder = Box::new(move |_, _: &_| Ok(handler.into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 1); assert_eq!(resp_vec[0].get_data(), [1, 2, 7]); assert_eq!(counter.load(atomic::Ordering::SeqCst), 0); // handler returns `None` but `finished == false` should not be called again. let counter = Arc::new(atomic::AtomicIsize::new(0)); let counter_clone = Arc::clone(&counter); let handler = StreamFromClosure::new(move |nth| match nth { 0 => { let mut resp = coppb::Response::default(); resp.set_data(vec![1, 2, 13]); Ok((Some(resp), false)) } 1 => Ok((None, false)), _ => { counter_clone.store(1, atomic::Ordering::SeqCst); Err(box_err!("unreachable")) } }); let handler_builder = Box::new(move |_, _: &_| Ok(handler.into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 1); assert_eq!(resp_vec[0].get_data(), [1, 2, 13]); assert_eq!(counter.load(atomic::Ordering::SeqCst), 0); // handler returns `Err(..)` should not be called again. let counter = Arc::new(atomic::AtomicIsize::new(0)); let counter_clone = Arc::clone(&counter); let handler = StreamFromClosure::new(move |nth| match nth { 0 => { let mut resp = coppb::Response::default(); resp.set_data(vec![1, 2, 23]); Ok((Some(resp), false)) } 1 => Err(box_err!("foo")), _ => { counter_clone.store(1, atomic::Ordering::SeqCst); Err(box_err!("unreachable")) } }); let handler_builder = Box::new(move |_, _: &_| Ok(handler.into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 2); assert_eq!(resp_vec[0].get_data(), [1, 2, 23]); assert!(!resp_vec[1].get_other_error().is_empty()); assert_eq!(counter.load(atomic::Ordering::SeqCst), 0); } #[test] fn test_channel_size() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config { end_point_stream_channel_size: 3, ..Config::default() }, read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let counter = Arc::new(atomic::AtomicIsize::new(0)); let counter_clone = Arc::clone(&counter); let handler = StreamFromClosure::new(move |nth| { // produce an infinite stream let mut resp = coppb::Response::default(); resp.set_data(vec![1, 2, nth as u8]); counter_clone.fetch_add(1, atomic::Ordering::SeqCst); Ok((Some(resp), false)) }); let handler_builder = Box::new(move |_, _: &_| Ok(handler.into_boxed())); let resp_vec = block_on_stream( copr.handle_stream_request(ParseCopRequestResult::default_for_test(handler_builder)) .unwrap(), ) .take(7) .collect::>>() .unwrap(); assert_eq!(resp_vec.len(), 7); assert!(counter.load(atomic::Ordering::SeqCst) < 14); } #[test] fn test_handle_time() { use tikv_util::config::ReadableDuration; /// Asserted that the snapshot can be retrieved in 500ms. const SNAPSHOT_DURATION_MS: u64 = 500; /// Asserted that the delay caused by OS scheduling other tasks is /// smaller than 200ms. This is mostly for CI. const HANDLE_ERROR_MS: u64 = 200; /// The acceptable error range for a coarse timer. Note that we use /// CLOCK_MONOTONIC_COARSE which can be slewed by time /// adjustment code (e.g., NTP, PTP). const COARSE_ERROR_MS: u64 = 50; /// The duration that payload executes. const PAYLOAD_SMALL: u64 = 3000; const PAYLOAD_LARGE: u64 = 6000; let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig { low_concurrency: 1, normal_concurrency: 1, high_concurrency: 1, ..CoprReadPoolConfig::default_for_test() }, engine, )); let config = Config { end_point_request_max_handle_duration: ReadableDuration::millis( (PAYLOAD_SMALL + PAYLOAD_LARGE) * 2, ), ..Default::default() }; let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &config, read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let (tx, rx) = std::sync::mpsc::channel(); // A request that requests execution details. let req_with_exec_detail: ReqContext = { let mut inner = ReqContextInner::default_for_test(); inner.context.set_record_time_stat(true); inner.into() }; { let mut wait_time: u64 = 0; // Request 1: Unary, success response. let handler_builder = Box::new(|_, _: &_| { Ok( UnaryFixture::new_with_duration(Ok(coppb::Response::default()), PAYLOAD_SMALL) .into_boxed(), ) }); let resp_future_1 = copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_with_exec_detail.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, }); let sender = tx.clone(); thread::spawn(move || sender.send(vec![block_on(resp_future_1).unwrap()]).unwrap()); // Sleep a while to make sure that thread is spawn and snapshot is taken. thread::sleep(Duration::from_millis(SNAPSHOT_DURATION_MS)); // Request 2: Unary, error response. let handler_builder = Box::new(|_, _: &_| { Ok( UnaryFixture::new_with_duration(Err(box_err!("foo")), PAYLOAD_LARGE) .into_boxed(), ) }); let resp_future_2 = copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_with_exec_detail.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, }); let sender = tx.clone(); thread::spawn(move || sender.send(vec![block_on(resp_future_2).unwrap()]).unwrap()); thread::sleep(Duration::from_millis(SNAPSHOT_DURATION_MS)); // Response 1 let resp = &rx.recv().unwrap()[0]; assert!(resp.get_other_error().is_empty()); assert_ge!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL + HANDLE_ERROR_MS + COARSE_ERROR_MS ); assert_ge!( resp.get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time.saturating_sub(HANDLE_ERROR_MS + COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time + HANDLE_ERROR_MS + COARSE_ERROR_MS ); wait_time += PAYLOAD_SMALL - SNAPSHOT_DURATION_MS; // Response 2 let resp = &rx.recv().unwrap()[0]; assert!(!resp.get_other_error().is_empty()); assert_ge!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE + HANDLE_ERROR_MS + COARSE_ERROR_MS ); assert_ge!( resp.get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time.saturating_sub(HANDLE_ERROR_MS + COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time + HANDLE_ERROR_MS + COARSE_ERROR_MS ); } { // Test multi-stage tasks // Request 1: Unary, success response. let handler_builder = Box::new(|_, _: &_| { Ok(UnaryFixture::new_with_duration_yieldable( Ok(coppb::Response::default()), PAYLOAD_SMALL, ) .into_boxed()) }); let resp_future_1 = copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_with_exec_detail.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, }); let sender = tx.clone(); thread::spawn(move || sender.send(vec![block_on(resp_future_1).unwrap()]).unwrap()); // Sleep a while to make sure that thread is spawn and snapshot is taken. thread::sleep(Duration::from_millis(SNAPSHOT_DURATION_MS)); // Request 2: Unary, error response. let handler_builder = Box::new(|_, _: &_| { Ok( UnaryFixture::new_with_duration_yieldable(Err(box_err!("foo")), PAYLOAD_LARGE) .into_boxed(), ) }); let resp_future_2 = copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_with_exec_detail.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, }); let sender = tx.clone(); thread::spawn(move || sender.send(vec![block_on(resp_future_2).unwrap()]).unwrap()); thread::sleep(Duration::from_millis(SNAPSHOT_DURATION_MS)); // Response 1 // // Note: `process_wall_time_ms` includes `total_process_time` and // `total_suspend_time`. Someday it will be separated, but for now, // let's just consider the combination. // // In the worst case, `total_suspend_time` could be totally req2 payload. // So here: req1 payload <= process time <= (req1 payload + req2 payload) let resp = &rx.recv().unwrap()[0]; assert!(resp.get_other_error().is_empty()); assert_ge!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL + PAYLOAD_LARGE + HANDLE_ERROR_MS + COARSE_ERROR_MS ); // Response 2 // // Note: `process_wall_time_ms` includes `total_process_time` and // `total_suspend_time`. Someday it will be separated, but for now, // let's just consider the combination. // // In the worst case, `total_suspend_time` could be totally req1 payload. // So here: req2 payload <= process time <= (req1 payload + req2 payload) let resp = &rx.recv().unwrap()[0]; assert!(!resp.get_other_error().is_empty()); assert_ge!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL + PAYLOAD_LARGE + HANDLE_ERROR_MS + COARSE_ERROR_MS ); } { let mut wait_time: u64 = 0; // Request 1: Unary, success response. let handler_builder = Box::new(|_, _: &_| { Ok( UnaryFixture::new_with_duration(Ok(coppb::Response::default()), PAYLOAD_LARGE) .into_boxed(), ) }); let resp_future_1 = copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_with_exec_detail.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, }); let sender = tx.clone(); thread::spawn(move || sender.send(vec![block_on(resp_future_1).unwrap()]).unwrap()); // Sleep a while to make sure that thread is spawn and snapshot is taken. thread::sleep(Duration::from_millis(SNAPSHOT_DURATION_MS)); // Request 2: Stream. let handler_builder = Box::new(|_, _: &_| { Ok(StreamFixture::new_with_duration( vec![ Ok(coppb::Response::default()), Err(box_err!("foo")), Ok(coppb::Response::default()), ], vec![PAYLOAD_SMALL, PAYLOAD_LARGE, PAYLOAD_SMALL], ) .into_boxed()) }); let resp_future_3 = copr .handle_stream_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_with_exec_detail.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, }) .unwrap() .map(|x| x.map(|x| x.into())); thread::spawn(move || { tx.send( block_on_stream(resp_future_3) .collect::>>() .unwrap(), ) .unwrap() }); // Response 1 let resp = &rx.recv().unwrap()[0]; assert!(resp.get_other_error().is_empty()); assert_ge!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE + HANDLE_ERROR_MS + COARSE_ERROR_MS ); assert_ge!( resp.get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time.saturating_sub(HANDLE_ERROR_MS + COARSE_ERROR_MS) ); assert_lt!( resp.get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time + HANDLE_ERROR_MS + COARSE_ERROR_MS ); wait_time += PAYLOAD_LARGE - SNAPSHOT_DURATION_MS; // Response 2 let resp = &rx.recv().unwrap(); assert_eq!(resp.len(), 2); assert!(resp[0].get_other_error().is_empty()); assert_ge!( resp[0] .get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp[0] .get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_SMALL + HANDLE_ERROR_MS + COARSE_ERROR_MS ); assert_ge!( resp[0] .get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time.saturating_sub(HANDLE_ERROR_MS + COARSE_ERROR_MS) ); assert_lt!( resp[0] .get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time + HANDLE_ERROR_MS + COARSE_ERROR_MS ); assert!(!resp[1].get_other_error().is_empty()); assert_ge!( resp[1] .get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE.saturating_sub(COARSE_ERROR_MS) ); assert_lt!( resp[1] .get_exec_details() .get_time_detail() .get_process_wall_time_ms(), PAYLOAD_LARGE + HANDLE_ERROR_MS + COARSE_ERROR_MS ); assert_ge!( resp[1] .get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time.saturating_sub(HANDLE_ERROR_MS + COARSE_ERROR_MS) ); assert_lt!( resp[1] .get_exec_details() .get_time_detail() .get_wait_wall_time_ms(), wait_time + HANDLE_ERROR_MS + COARSE_ERROR_MS ); } } #[test] fn test_exceed_deadline() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); { let handler_builder = Box::new(|_, _: &_| { thread::sleep(Duration::from_millis(600)); Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed()) }); let mut inner = ReqContextInner::default_for_test(); inner.deadline = Deadline::from_now(Duration::from_millis(500)); let config: ReqContext = inner.into(); let resp = block_on(copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: config, remote_dispatch_precharge_bytes: 0, handler_builder, })) .unwrap(); assert_eq!(resp.get_data().len(), 0); let region_err = resp.get_region_error(); assert_eq!( region_err.get_server_is_busy().reason, "deadline is exceeded".to_string() ); } { let handler_builder = Box::new(|_, _: &_| { Ok( UnaryFixture::new_with_duration_yieldable(Ok(coppb::Response::default()), 1500) .into_boxed(), ) }); let mut inner = ReqContextInner::default_for_test(); inner.deadline = Deadline::from_now(Duration::from_millis(500)); let config: ReqContext = inner.into(); let resp = block_on(copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: config, remote_dispatch_precharge_bytes: 0, handler_builder, })) .unwrap(); assert_eq!(resp.get_data().len(), 0); let region_err = resp.get_region_error(); assert_eq!( region_err.get_server_is_busy().reason, "deadline is exceeded".to_string() ); } } #[test] fn test_check_memory_locks() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let key = Key::from_raw(b"key"); let guard = block_on(cm.lock_key(&key)); guard.with_lock(|lock| { *lock = Some(txn_types::Lock::new( LockType::Put, b"key".to_vec(), 10.into(), 100, Some(vec![]), 0.into(), 1, 20.into(), )); }); let config = Config::default(); let copr = Endpoint::::new( &config, read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let mut req = coppb::Request::default(); req.mut_context().set_isolation_level(IsolationLevel::Si); req.set_start_ts(100); req.set_tp(REQ_TYPE_DAG); let mut key_range = coppb::KeyRange::default(); key_range.set_start(b"a".to_vec()); key_range.set_end(b"z".to_vec()); req.mut_ranges().push(key_range); let mut dag = DagRequest::default(); dag.mut_executors().push(Executor::default()); req.set_data(dag.write_to_bytes().unwrap()); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert_eq!(resp.get_locked().get_key(), b"key"); } #[test] fn test_make_error_response() { let resp = make_error_response(Error::DeadlineExceeded); let region_err = resp.get_region_error(); assert_eq!( region_err.get_server_is_busy().reason, "deadline is exceeded".to_string() ); assert_eq!( region_err.get_message(), "Coprocessor task terminated due to exceeding the deadline" ); } #[test] fn test_make_error_response_cpu_throttled() { let msg = "[CPU throttle] Keyspace 1 request timeout during CPU token allocation"; let resp = make_error_response(Error::CpuThrottled(msg.to_string())); let region_err = resp.get_region_error(); assert!(region_err.has_server_is_busy()); assert!( region_err .get_server_is_busy() .reason .contains("CPU throttled"), "reason should contain CPU throttle info, got: {}", region_err.get_server_is_busy().reason ); assert!( region_err.get_message().contains(msg), "message should contain original error, got: {}", region_err.get_message() ); } #[test] fn test_low_priority_dag_keyspace_wait_respects_deadline() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.low_priority_dag_global_enabled = true; cpu_throttle_config.low_priority_dag_global_ratio = 1.0; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = false; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 10_000, std::collections::HashMap::new(), CoprReadPoolConfig { low_concurrency: 2, ..CoprReadPoolConfig::default_for_test() }, ); let keyspace_read_limiter = KeyspaceReadLimiter::new(1, Duration::from_secs(1), Duration::from_secs(5)); keyspace_read_limiter.set_speed_limit(Some(1.0), None); keyspace_read_limiter.consume_request(); let req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Low, Duration::from_millis(50)); let handler = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let start = std::time::Instant::now(); let err = run_tracked_unary_request_with_limiters( copr, engine, req_ctx, handler, None, Some(keyspace_read_limiter), ) .unwrap_err(); assert!(matches!(err, Error::DeadlineExceeded)); assert_lt!(start.elapsed(), Duration::from_millis(300)); } #[test] fn test_low_priority_dag_cpu_throttled_by_global_bucket() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.low_priority_dag_global_enabled = true; cpu_throttle_config.low_priority_dag_global_ratio = 1.0; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = false; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 1_000, std::collections::HashMap::new(), CoprReadPoolConfig { low_concurrency: 2, ..CoprReadPoolConfig::default_for_test() }, ); let blocking_req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Low, Duration::from_secs(2)); let (started_tx, started_rx) = mpsc::channel(); let blocking_handler = Box::new(move |_, _: &_| { Ok(UnaryFixture::new_with_duration_and_start_signal( Ok(coppb::Response::default()), 1_000, started_tx, ) .into_boxed()) }); let (tx, rx) = mpsc::channel(); let copr_clone = copr.clone(); let engine_clone = engine.clone(); thread::spawn(move || { tx.send(run_tracked_unary_request( copr_clone, engine_clone, blocking_req_ctx, blocking_handler, )) .unwrap(); }); started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); wait_for_low_priority_bucket_available(&copr, 0); let throttled_req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Low, Duration::from_millis(500)); let throttled_handler = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let err = run_tracked_unary_request(copr, engine, throttled_req_ctx, throttled_handler) .unwrap_err(); match err { Error::CpuThrottled(msg) => { assert!( msg.contains("Low priority DAG global CPU quota exhausted"), "unexpected message: {}", msg ); } other => panic!("unexpected error: {:?}", other), } let first_resp = rx.recv().unwrap().unwrap(); assert!(first_resp.get_other_error().is_empty()); } #[test] fn test_fair_global_cpu_gate_skips_dag_cpu_admission() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = false; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 1_000, std::collections::HashMap::new(), CoprReadPoolConfig::default_for_test(), ); let manager = copr .read_limiter .as_ref() .and_then(|read_limiter| read_limiter.get_cpu_throttle_manager()) .expect("cpu throttle manager should exist"); assert!(!manager.consume_fair_global_cpu(1_000)); assert_eq!(manager.global_available(), 0); let req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Normal, Duration::from_secs(2)); let handler = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let resp = run_tracked_unary_request_with_fair_gate(copr, engine, req_ctx, handler, true, None) .unwrap(); assert!(resp.get_other_error().is_empty()); assert_eq!(manager.global_available(), 0); } #[test] fn test_fair_global_cpu_gate_does_not_install_legacy_runtime_budget() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = true; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 10_000, std::collections::HashMap::new(), CoprReadPoolConfig::default_for_test(), ); let read_limiter = copr.read_limiter.clone().unwrap(); let poll_cpu_recorder = Arc::new(CoprReadPoolPollCpuRecorder::new( 1, 0, false, copr.resource_publisher.clone().unwrap(), )); let runtime_budget_installed = Arc::new(AtomicBool::new(false)); let runtime_budget_installed_clone = runtime_budget_installed.clone(); let req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Normal, Duration::from_secs(2)); let handler = Box::new(move |_, _: &_| { Ok(UnaryFixture::new_with_runtime_budget_observer( Ok(coppb::Response::default()), runtime_budget_installed_clone.clone(), ) .into_boxed()) }); let resp = run_tracked_unary_request_with_fair_gate_and_recorder( copr, engine, req_ctx, handler, true, Some(read_limiter), None, Some(poll_cpu_recorder), ) .unwrap(); assert!(resp.get_other_error().is_empty()); assert!(!runtime_budget_installed.load(Ordering::SeqCst)); } #[test] fn test_fair_global_cpu_gate_skips_remote_fallback_dag_cpu_admission() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = false; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, _engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 1_000, std::collections::HashMap::new(), CoprReadPoolConfig::default_for_test(), ); let read_limiter = copr.read_limiter.clone().unwrap(); let manager = read_limiter .get_cpu_throttle_manager() .expect("cpu throttle manager should exist"); assert!(!manager.consume_fair_global_cpu(1_000)); assert_eq!(manager.global_available(), 0); let req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Normal, Duration::from_secs(2)); let fallback_admission = Endpoint::::build_dag_fallback_admission( &req_ctx, Some(read_limiter), copr.resource_publisher.clone(), true, false, false, true, ) .expect("resource control should build fallback admission"); let admission = block_on(fallback_admission()).unwrap(); assert!(admission.is_none()); assert_eq!(manager.global_available(), 0); } #[test] fn test_fairness_construction_gate_skips_local_and_fallback_legacy_dag_admission() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = true; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config.clone(), 1_000, std::collections::HashMap::new(), CoprReadPoolConfig::default_for_test(), ); let mut resource_control_config = resource_control::Config::default(); resource_control_config.enabled = true; let read_limiter = ReadLimiter::new_with_legacy_read_limiter(resource_control_config, false); let manager = Arc::new(CpuThrottleManager::new( cpu_throttle_config, 1_000, std::collections::HashMap::new(), )); read_limiter.set_cpu_throttle_manager(manager.clone()); assert!(!manager.consume_fair_global_cpu(1_000)); assert_eq!(manager.global_available(), 0); let req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Normal, Duration::from_secs(2)); assert!( Endpoint::::build_dag_fallback_admission( &req_ctx, Some(read_limiter.clone()), copr.resource_publisher.clone(), true, false, false, false, ) .is_none() ); let handler = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let resp = run_tracked_unary_request_with_fair_gate( copr, engine, req_ctx, handler, false, Some(read_limiter), ) .unwrap(); assert!(resp.get_other_error().is_empty()); assert_eq!(manager.global_available(), 0); } #[test] fn test_normal_priority_dag_not_blocked_by_low_priority_bucket() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.low_priority_dag_global_enabled = true; cpu_throttle_config.low_priority_dag_global_ratio = 1.0; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = false; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 1_000, std::collections::HashMap::new(), CoprReadPoolConfig { low_concurrency: 2, normal_concurrency: 2, ..CoprReadPoolConfig::default_for_test() }, ); let low_req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Low, Duration::from_secs(2)); let (started_tx, started_rx) = mpsc::channel(); let low_handler = Box::new(move |_, _: &_| { Ok(UnaryFixture::new_with_duration_and_start_signal( Ok(coppb::Response::default()), 1_000, started_tx, ) .into_boxed()) }); let (tx, rx) = mpsc::channel(); let copr_clone = copr.clone(); let engine_clone = engine.clone(); thread::spawn(move || { tx.send(run_tracked_unary_request( copr_clone, engine_clone, low_req_ctx, low_handler, )) .unwrap(); }); started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); wait_for_low_priority_bucket_available(&copr, 0); let normal_req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Normal, Duration::from_millis(500)); let normal_handler = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let resp = run_tracked_unary_request(copr, engine, normal_req_ctx, normal_handler).unwrap(); assert!( !resp.has_region_error(), "normal priority request should bypass low-priority bucket: {:?}", resp.get_region_error() ); assert!(resp.get_other_error().is_empty()); let low_resp = rx.recv().unwrap().unwrap(); assert!(low_resp.get_other_error().is_empty()); } #[test] fn test_low_priority_batch_dag_throttled_by_global_bucket() { let mut cpu_throttle_config = resource_control::CpuThrottleConfig::default(); cpu_throttle_config.enabled = true; cpu_throttle_config.low_priority_dag_global_enabled = true; cpu_throttle_config.low_priority_dag_global_ratio = 1.0; cpu_throttle_config.estimated_cpu_per_request_us = 1_000; cpu_throttle_config.enable_runtime_token_management = false; cpu_throttle_config.enable_burst = false; cpu_throttle_config.refill_interval_ms = 10_000; let mut keyspace_cpu_us = std::collections::HashMap::new(); keyspace_cpu_us.insert(1, 2_000); let (copr, engine) = build_endpoint_with_cpu_throttle( cpu_throttle_config, 1_000, keyspace_cpu_us, CoprReadPoolConfig { low_concurrency: 2, normal_concurrency: 2, ..CoprReadPoolConfig::default_for_test() }, ); let blocking_req_ctx = make_req_ctx_with_priority(1, kvrpcpb::CommandPri::Low, Duration::from_secs(2)); let (started_tx, started_rx) = mpsc::channel(); let blocking_handler = Box::new(move |_, _: &_| { Ok(UnaryFixture::new_with_duration_and_start_signal( Ok(coppb::Response::default()), 1_000, started_tx, ) .into_boxed()) }); let (tx, rx) = mpsc::channel(); let copr_clone = copr.clone(); let engine_clone = engine.clone(); thread::spawn(move || { tx.send(run_tracked_unary_request( copr_clone, engine_clone, blocking_req_ctx, blocking_handler, )) .unwrap(); }); started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); wait_for_low_priority_bucket_available(&copr, 0); let req = make_batch_dag_request(1, kvrpcpb::CommandPri::Low, 200); let resp = block_on(copr.parse_and_handle_unary_request(req, None)); assert!( resp.get_region_error() .get_server_is_busy() .reason .contains("Low priority DAG global CPU quota exhausted"), "main request should be throttled by low-priority bucket: {:?}", resp.get_region_error() ); assert_eq!(resp.get_batch_responses().len(), 1); let batch_resp = &resp.get_batch_responses()[0]; assert!( batch_resp .get_region_error() .get_server_is_busy() .reason .contains("Low priority DAG global CPU quota exhausted"), "batch request should be throttled by low-priority bucket: {:?}", batch_resp.get_region_error() ); let blocking_resp = rx.recv().unwrap().unwrap(); assert!(blocking_resp.get_other_error().is_empty()); } #[test] fn test_make_error_delegate_response_cpu_throttled() { let msg = "[CPU throttle] Request timeout waiting for CPU tokens"; let resp = make_error_delegate_response(Error::CpuThrottled(msg.to_string())); let region_err = resp.get_region_error(); assert!(region_err.has_server_is_busy()); assert!( region_err .get_server_is_busy() .reason .contains("CPU throttled"), "reason should contain CPU throttle info, got: {}", region_err.get_server_is_busy().reason ); assert!( region_err.get_message().contains(msg), "message should contain original error, got: {}", region_err.get_message() ); } // A default ReqContext that support to access another snapshot in a // request. It can be used to create an not None value of // Option fn default_req_ctx_support_snap_accessor() -> ReqContextInner { let mut pb_ctx = kvrpcpb::Context::default(); pb_ctx.set_isolation_level(IsolationLevel::Si); pb_ctx.set_stale_read(false); pb_ctx.set_replica_read(false); pb_ctx.set_peer(metapb::Peer { id: 12, store_id: 100, ..Default::default() }); pb_ctx.set_region_id(23); let mut req_ctx = ReqContextInner::default_for_test(); req_ctx.context = pb_ctx; req_ctx.txn_start_ts = TimeStamp::new(1234567); req_ctx.is_desc_scan = Some(true); req_ctx.ranges = vec![ coppb::KeyRange { start: b"a1".to_vec(), end: b"b1".to_vec(), ..Default::default() }, coppb::KeyRange { start: b"b3".to_vec(), end: b"b4".to_vec(), ..Default::default() }, ]; req_ctx } #[test] fn test_secondary_snap_store_accessor_new() { type StoreAccessor = ExtraSnapStoreAccessor; // construct a ReqContext that support to access another snapshot in a request let req_ctx = default_req_ctx_support_snap_accessor(); // accessor support case let ri_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![]); let mut ctx = req_ctx.clone(); assert!( StoreAccessor::new( ctx.into(), Some(ri_accessor.clone()), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .is_some() ); // does not support Rc / RcCheckTs ctx = req_ctx.clone(); ctx.context.set_isolation_level(IsolationLevel::Rc); assert!( StoreAccessor::new( ctx.into(), Some(ri_accessor.clone()), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .is_none() ); ctx = req_ctx.clone(); ctx.context.set_isolation_level(IsolationLevel::RcCheckTs); assert!( StoreAccessor::new( ctx.into(), Some(ri_accessor.clone()), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .is_none() ); // does not support stale read ctx = req_ctx.clone(); ctx.context.set_stale_read(true); assert!( StoreAccessor::new( ctx.into(), Some(ri_accessor.clone()), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .is_none() ); // does not support replica read ctx = req_ctx.clone(); ctx.context.set_replica_read(true); assert!( StoreAccessor::new( ctx.into(), Some(ri_accessor.clone()), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .is_none() ); } #[test] fn test_extra_snap_store_accessor_locate_region_by_key() { set_tls_engine(TestEngineBuilder::new().build().unwrap()); defer! { unsafe {destroy_tls_engine::()} } let (r1, r2, r3) = ( metapb::Region { id: 1, start_key: b"b".to_vec(), end_key: b"d".to_vec(), ..Default::default() }, metapb::Region { id: 2, start_key: b"e".to_vec(), end_key: b"g".to_vec(), ..Default::default() }, metapb::Region { id: 3, start_key: b"g".to_vec(), end_key: b"h".to_vec(), ..Default::default() }, ); let ri_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![ RegionInfo::new(r1.clone(), StateRole::Leader, 0), RegionInfo::new(r2.clone(), StateRole::Follower, 0), RegionInfo::new(r3.clone(), StateRole::Leader, 0), ]); type StoreAccessor = ExtraSnapStoreAccessor; let accessor = StoreAccessor::new( default_req_ctx_support_snap_accessor().into(), Some(ri_accessor), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .unwrap(); // key is before any region, not found assert_eq!( block_on(accessor.find_region_by_key(b"a")), FindRegionResult::with_not_found(Some(b"b".to_vec())), ); // key is region start, found assert_eq!( block_on(accessor.find_region_by_key(b"b")), FindRegionResult::with_found(r1.clone(), StateRole::Leader), ); // key is in a region range, found assert_eq!( block_on(accessor.find_region_by_key(b"b1")), FindRegionResult::with_found(r1.clone(), StateRole::Leader), ); assert_eq!( block_on(accessor.find_region_by_key(b"c")), FindRegionResult::with_found(r1.clone(), StateRole::Leader), ); assert_eq!( block_on(accessor.find_region_by_key(b"c1")), FindRegionResult::with_found(r1.clone(), StateRole::Leader), ); // key is a region's end, but not any region's end, not found assert_eq!( block_on(accessor.find_region_by_key(b"d")), FindRegionResult::with_not_found(Some(b"e".to_vec())), ); // key is between a region's end and another region's start, not found assert_eq!( block_on(accessor.find_region_by_key(b"d1")), FindRegionResult::with_not_found(Some(b"e".to_vec())), ); // key is in a region, but the region is not a leader assert_eq!( block_on(accessor.find_region_by_key(b"e")), FindRegionResult::with_found(r2.clone(), StateRole::Follower), ); assert_eq!( block_on(accessor.find_region_by_key(b"f")), FindRegionResult::with_found(r2.clone(), StateRole::Follower), ); // key is a region's end, and another region's start found assert_eq!( block_on(accessor.find_region_by_key(b"g")), FindRegionResult::with_found(r3.clone(), StateRole::Leader), ); // key is after all regions, not found without next_region_start assert_eq!( block_on(accessor.find_region_by_key(b"h")), FindRegionResult::with_not_found(None), ); assert_eq!( block_on(accessor.find_region_by_key(b"i")), FindRegionResult::with_not_found(None), ); } #[test] fn test_secondary_snap_store_accessor_get_local_region_storage() { type StoreAccessor = ExtraSnapStoreAccessor; #[derive(Clone)] struct TestCtx { store_id: u64, req_ctx: Arc>, req_region: Arc>, req_ranges: Arc>>, called: Arc, } impl TestCtx { fn new_accessor(&self) -> StoreAccessor { let region = self.req_region.lock().unwrap().clone(); let region = RegionInfo::new(region, StateRole::Leader, 0); let ri_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![region]); StoreAccessor::new( self.get_req_ctx(), Some(ri_accessor), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .unwrap() } fn get_req_ctx(&self) -> ReqContext { ReqContext(Arc::new(self.req_ctx.lock().unwrap().clone())) } fn get_req_region(&self) -> metapb::Region { self.req_region.lock().unwrap().clone() } fn get_req_ranges(&self) -> Vec { self.req_ranges.lock().unwrap().clone() } fn get_local_region_storage_with_check( &self, ) -> StorageResult::Snap>> { let accessor = &self.new_accessor(); assert!(!self.called.load(atomic::Ordering::SeqCst)); let result = block_on( accessor .get_local_region_storage(&self.get_req_region(), &self.get_req_ranges()), ); let called = self.called.swap(false, atomic::Ordering::SeqCst); if let Ok(ref store) = result { assert!(called); let req_ctx = self.get_req_ctx(); let pb_ctx = &req_ctx.context; assert_eq!(store.get_start_ts(), req_ctx.txn_start_ts.into_inner()); assert_eq!(store.is_fill_cache(), !pb_ctx.get_not_fill_cache()); assert_eq!(store.get_by_pass_locks(), req_ctx.bypass_locks); // check_has_newer_ts_data should always be false to avoid caching the cop // response assert!(!store.is_check_has_newer_ts_data()); } result } } let test_ctx = { let req_ctx = Arc::new(Mutex::new(default_req_ctx_support_snap_accessor())); let store_id = req_ctx.lock().unwrap().context.get_peer().get_store_id(); assert!(store_id > 0); TestCtx { store_id, req_ctx, req_region: Arc::new(Mutex::new(metapb::Region { id: 123, start_key: b"a".to_vec(), end_key: b"z".to_vec(), peers: vec![ metapb::Peer { id: 1, store_id: store_id - 1, ..Default::default() }, metapb::Peer { id: 2, store_id, ..Default::default() }, metapb::Peer { id: 3, store_id: store_id + 1, ..Default::default() }, ] .into(), ..Default::default() })), req_ranges: Arc::new(Mutex::new(vec![coppb::KeyRange { start: b"a".to_vec(), end: b"b".to_vec(), ..Default::default() }])), called: Arc::new(atomic::AtomicBool::new(false)), } }; let test_ctx_for_engine = test_ctx.clone(); set_tls_engine( MockEngineBuilder::from_engine(TestEngineBuilder::new().build().unwrap()) // set the hook to check the snap_ctx as the argument of Engine::async_snapshot .set_pre_async_snapshot(move |snap_ctx| { let test_ctx = test_ctx_for_engine.clone(); assert!(!test_ctx.called.swap( true, atomic::Ordering::SeqCst, )); let check_ctx = test_ctx.get_req_ctx(); let check_region = test_ctx.get_req_region(); // Currently, only leader read is supported in this test. assert!(!check_ctx.context.get_replica_read() && !check_ctx.context.get_stale_read()); assert!(!check_ctx.ranges.is_empty()); // snap_ctx.pb_ctx should be the same as ReqContext.context assert_eq!(snap_ctx.pb_ctx.clone(), check_ctx.context.clone()); // snap_ctx.extra_snap_override should be present with the correct region info assert_eq!( snap_ctx.extra_region_override, Some(ExtraRegionOverride { region_id: check_region.id, region_epoch: check_region.get_region_epoch().clone(), peer: check_region.get_peers()[1].clone(), check_term: None, }) ); // should select a peer with the right store_id. assert_eq!( snap_ctx.extra_region_override.as_ref().unwrap().peer.store_id, test_ctx.store_id ); // snapshot cache is not supported currently, so snap_ctx.read_id is always None. assert!(snap_ctx.read_id.is_none()); // snap_ctx.start_ts should be the same as req_ctx.txn_start_ts assert!(check_ctx.txn_start_ts > TimeStamp::zero()); assert_eq!(snap_ctx.start_ts, Some(check_ctx.txn_start_ts)); // Even if req_ctx.ranges is not empty, snap_ctx.key_ranges should be empty in // leader read because leader read does not need to check keys locks. assert!(!test_ctx.get_req_ranges().is_empty()); assert_eq!(snap_ctx.key_ranges.len(), 0); }) .build(), ); defer! { unsafe {destroy_tls_engine::()} } // normal case test_ctx .get_local_region_storage_with_check() .expect("should succeed"); // if set_not_fill_cache, it should work { let mut req_ctx = test_ctx.req_ctx.lock().unwrap(); // in previous test, get_not_fill_cache should return false. assert!(!req_ctx.context.get_not_fill_cache()); req_ctx.context.set_not_fill_cache(true); } test_ctx .get_local_region_storage_with_check() .expect("should succeed"); // cannot find peer { test_ctx.req_region.lock().unwrap().peers[1].store_id = 99999999; } let err = test_ctx .get_local_region_storage_with_check() .err() .unwrap(); assert!(err.to_string().contains("cannot find peer in region")); { test_ctx.req_region.lock().unwrap().peers[1].store_id = test_ctx.store_id; } // should return error when engine returns error unsafe { with_tls_engine(|e: &mut MockEngine| e.base_engine().trigger_not_leader()); } let err = test_ctx .get_local_region_storage_with_check() .err() .unwrap(); assert!(err.to_string().contains("not_leader")); } #[test] #[ignore = "CSE incompatibility; see #6158"] fn test_secondary_tikv_storage_accessor() { set_tls_engine(TestEngineBuilder::new().build().unwrap()); defer! { unsafe {destroy_tls_engine::()} } let def_req = default_req_ctx_support_snap_accessor(); let store_id = def_req.context.get_peer().get_store_id(); let ri_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![]); let store_accessor = ExtraSnapStoreAccessor::::new( def_req.into(), Some(ri_accessor), ConcurrencyManager::new(1.into()), ReadRescheduleOptions::default(), ) .unwrap(); let storage_accessor = dag::ExtraStorageAccessor::>::from_store_accessor( store_accessor, ); let result = block_on( storage_accessor.get_local_region_storage( &metapb::Region { id: 123, start_key: b"a".to_vec(), end_key: b"z".to_vec(), peers: vec![metapb::Peer { id: 1, store_id, ..Default::default() }] .into(), ..Default::default() }, &[coppb::KeyRange { start: b"a".to_vec(), end: b"b".to_vec(), ..Default::default() }], ), ); assert!(result.is_err()); let err = result.err().unwrap(); assert!(err.to_string().contains("unexpected snapshot")); } #[test] fn test_extra_snap_accessor_check_memory_locks() { let engine = TestEngineBuilder::new().build().unwrap(); set_tls_engine(engine); defer! { unsafe {destroy_tls_engine::()} } let region = metapb::Region { id: 1, start_key: b"".to_vec(), end_key: b"".to_vec(), peers: vec![metapb::Peer { id: 1, store_id: 100, ..Default::default() }] .into(), ..Default::default() }; let cm = ConcurrencyManager::new(1.into()); let mut req = default_req_ctx_support_snap_accessor(); req.txn_start_ts = 100.into(); let ri_accessor = RegionInfoAccessor::new_with_regions_for_test(vec![RegionInfo::new( region.clone(), StateRole::Leader, 0, )]); let accessor = ExtraSnapStoreAccessor::::new( req.into(), Some(ri_accessor), cm.clone(), ReadRescheduleOptions::default(), ) .unwrap(); let key = Key::from_raw(b"key"); let guard = block_on(cm.lock_key(&key)); guard.with_lock(|lock| { *lock = Some(txn_types::Lock::new( LockType::Put, b"key".to_vec(), 10.into(), 100, Some(vec![]), 0.into(), 1, 20.into(), )); }); let err = block_on(accessor.get_local_region_storage( ®ion, &[coppb::KeyRange { start: b"key".to_vec(), end: b"key0".to_vec(), ..Default::default() }], )) .map_err(Error::from) .err() .unwrap(); assert_matches!(err, Error::Locked(LockInfo { key, .. }) if { assert_eq!(key, b"key".to_vec()); true }); } #[test] fn test_memory_quota() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); // By default, coprocessor does not return memory quota exceeded error. { let handler_builder = Box::new(|_, _: &_| { Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed()) }); let r = ParseCopRequestResult::default_for_test(handler_builder); let resp = block_on(copr.handle_unary_request(r)).unwrap(); assert!(!resp.has_region_error(), "{:?}", resp); } // Trigger memory quota exceeded error. copr.memory_quota.set_capacity(1); { let handler_builder = Box::new(|_, _: &_| { Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed()) }); let r = ParseCopRequestResult::default_for_test(handler_builder); let res = block_on(copr.handle_unary_request(r)); assert!(res.is_err(), "{:?}", res); let resp = make_error_response(res.unwrap_err()); assert_eq!(resp.get_data().len(), 0); let region_err = resp.get_region_error(); assert!( region_err .get_server_is_busy() .reason .contains("exceeding memory quota"), "{:?}", region_err.get_server_is_busy().reason ); } } // End-to-end: drive `end-point-memory-quota` through the real // CopConfigManager (the one Endpoint hands out), then exercise both the // admit-and-release and the reject paths. Guards the full online-config // wiring plus the OwnedAllocated drop contract (in_use must return to 0 // after every request, on either path). #[test] fn test_memory_quota_via_online_config() { use std::collections::HashMap; use online_config::{ConfigChange, ConfigValue}; let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let make_req = || { let handler_builder = Box::new(|_, _: &_| { Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed()) }); ParseCopRequestResult::default_for_test(handler_builder) }; // Default quota admits the request and releases it on completion. let resp = block_on(copr.handle_unary_request(make_req())).unwrap(); assert!(!resp.has_region_error(), "{:?}", resp); assert_eq!( copr.memory_quota.in_use(), 0, "OwnedAllocated must release bytes on success" ); // Drive a 1-byte quota through the online-config dispatch path. let mut change: ConfigChange = HashMap::new(); change.insert("end_point_memory_quota".to_owned(), ConfigValue::Size(1)); let mut cfg_mgr = copr.config_manager(); cfg_mgr.dispatch(change).unwrap(); assert_eq!(copr.memory_quota.capacity(), 1); // Rejection path: `ServerIsBusy` with "exceeding memory quota". let res = block_on(copr.handle_unary_request(make_req())); let resp = make_error_response(res.unwrap_err()); assert_eq!(resp.get_data().len(), 0); let region_err = resp.get_region_error(); assert!( region_err .get_server_is_busy() .reason .contains("exceeding memory quota"), "{:?}", region_err.get_server_is_busy().reason ); assert_eq!( copr.memory_quota.in_use(), 0, "rejected admission must not leave bytes accounted" ); // Raising the cap via the same online-config path resumes traffic. let mut change: ConfigChange = HashMap::new(); change.insert( "end_point_memory_quota".to_owned(), ConfigValue::Size(tikv_util::config::ReadableSize::mb(64).0), ); cfg_mgr.dispatch(change).unwrap(); assert_eq!( copr.memory_quota.capacity(), tikv_util::config::ReadableSize::mb(64).0 as usize ); let resp = block_on(copr.handle_unary_request(make_req())).unwrap(); assert!(!resp.has_region_error(), "{:?}", resp); assert_eq!(copr.memory_quota.in_use(), 0); } #[test] fn test_memory_quota_counts_remote_dispatch_precharge() { use std::sync::{ Arc, atomic::{AtomicBool, Ordering}, }; let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); let req_ctx = ReqContext::default_for_test(); copr.memory_quota .set_capacity(ReadableSize::mb(1).0 as usize); let handler_builder = Box::new(|_, _: &_| Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed())); let resp = block_on(copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx: req_ctx.clone(), remote_dispatch_precharge_bytes: 0, handler_builder, })) .unwrap(); assert!(resp.get_other_error().is_empty(), "{:?}", resp); assert!(!resp.has_region_error(), "{:?}", resp); let executed = Arc::new(AtomicBool::new(false)); let executed_for_builder = Arc::clone(&executed); let handler_builder = Box::new(move |_, _: &_| { executed_for_builder.store(true, Ordering::SeqCst); Ok(UnaryFixture::new(Ok(coppb::Response::default())).into_boxed()) }); let res = block_on(copr.handle_unary_request(ParseCopRequestResult { req_tag: ReqTag::test, req_ctx, remote_dispatch_precharge_bytes: ReadableSize::mb(8).0 as usize, handler_builder, })); assert!(res.is_err(), "{:?}", res); assert!( matches!(res, Err(Error::MemoryQuotaExceeded)), "expected MemoryQuotaExceeded, got {:?}", res ); assert!( !executed.load(Ordering::SeqCst), "handler must not execute after precharge quota rejection" ); assert_eq!(copr.memory_quota.in_use(), 0); } // Admission rejection on the stream path. `test_memory_quota` only covers // the unary path, so a regression in `handle_stream_request`'s // `read_pool_spawn_with_memory_quota_check` call would otherwise slip // through. #[test] fn test_memory_quota_stream_path() { let engine = TestEngineBuilder::new().build().unwrap(); let read_pool = ReadPool::from(build_read_pool_for_test( &CoprReadPoolConfig::default_for_test(), engine, )); let cm = ConcurrencyManager::new(1.into()); let copr = Endpoint::::new( &Config::default(), read_pool.handle(), cm, None, ResourceTagFactory::new_for_test(), Arc::new(SecurityManager::default()), ); copr.memory_quota.set_capacity(1); let handler_builder = Box::new(|_, _: &_| { Ok(StreamFixture::new(vec![Ok(coppb::Response::default())]).into_boxed()) }); let r = ParseCopRequestResult::default_for_test(handler_builder); let res = copr.handle_stream_request(r); assert!( matches!(res, Err(Error::MemoryQuotaExceeded)), "stream admission should reject with MemoryQuotaExceeded, got {:?}", res.err() ); assert_eq!(copr.memory_quota.in_use(), 0); } #[test] fn test_make_error_batch_response_cpu_throttled() { let msg = "[CPU throttle] Low priority DAG global CPU quota exhausted"; let mut resp = coppb::StoreBatchTaskResponse::new(); make_error_batch_response(&mut resp, Error::CpuThrottled(msg.to_string())); let region_err = resp.get_region_error(); assert!(region_err.has_server_is_busy()); assert!( region_err.get_server_is_busy().reason.contains(msg), "reason should contain original CPU throttle info, got: {}", region_err.get_server_is_busy().reason ); assert!( region_err.get_message().contains(msg), "message should contain original CPU throttle info, got: {}", region_err.get_message() ); } #[test] fn test_store_batch_request_uses_child_buckets_version() { let mut req = coppb::Request::default(); req.mut_context().set_buckets_version(7); let mut task = coppb::StoreBatchTask::new(); task.set_task_id(42); task.set_buckets_version(99); let (child_req, task_id) = Endpoint::::make_store_batch_request(&req, &mut task); assert_eq!(task_id, 42); assert_eq!(child_req.get_context().get_buckets_version(), 99); } #[test] fn test_store_batch_request_keeps_parent_buckets_version_when_child_zero() { let mut req = coppb::Request::default(); req.mut_context().set_buckets_version(7); let mut task = coppb::StoreBatchTask::new(); task.set_task_id(42); let (child_req, task_id) = Endpoint::::make_store_batch_request(&req, &mut task); assert_eq!(task_id, 42); assert_eq!(child_req.get_context().get_buckets_version(), 7); } }