From 275e45392eb03643fa988e459ca72a8af42bf833 Mon Sep 17 00:00:00 2001 From: zyguan Date: Sat, 10 Oct 2026 10:17:22 +0000 Subject: [PATCH] *: protect transaction RPC compatibility Signed-off-by: zyguan --- proto/VERSION | 6 +- proto/errorpb.proto | 32 ++ proto/kvrpcpb.proto | 41 +- proto/metapb.proto | 12 + src/common/errors.rs | 52 +- src/common/mod.rs | 4 + src/config.rs | 83 +++ src/generated/errorpb.rs | 63 +++ src/generated/kvrpcpb.rs | 72 ++- src/generated/metapb.rs | 16 + src/lib.rs | 4 + src/mock.rs | 94 ++++ src/pd/client.rs | 109 +++- src/raw/requests.rs | 14 + src/region_cache.rs | 46 +- src/request/mod.rs | 230 ++++++++- src/request/plan.rs | 678 +++++++++++++++++++++---- src/request/plan_builder.rs | 141 ++++- src/request/shard.rs | 39 ++ src/store/errors.rs | 37 ++ src/store/mod.rs | 146 +++++- src/store/request.rs | 490 ++++++++++++++++-- src/transaction/client.rs | 75 ++- src/transaction/lock.rs | 300 +++++++++-- src/transaction/mod.rs | 2 +- src/transaction/requests.rs | 96 +++- src/transaction/transaction.rs | 358 +++++++++++-- src/transaction/transaction/scanner.rs | 188 ++++++- 28 files changed, 3105 insertions(+), 323 deletions(-) diff --git a/proto/VERSION b/proto/VERSION index 6af27a69..a2edffc7 100644 --- a/proto/VERSION +++ b/proto/VERSION @@ -1,3 +1,3 @@ -pingcap/kvproto b41e86365ce0e0fba482005c0dfa5e2ca967d74c (2026-06-22) -Matches client-go v2.0.8-0.20260708122311-01bd8f99f4da (go.mod: kvproto v0.0.0-20260622063236-b41e86365ce0). -Re-vendor: copy proto/*.proto (matching names only) and include/ from that revision, then `make generate`. +pingcap/kvproto d80b6c5e7821b55f78092caead501ab0f00102da (2026-09-10) +Includes transaction RPC compatibility declarations and structured errors from kvproto PR #1534. +Re-vendor: copy the vendored proto files and any newly imported dependencies from that revision, then `make generate`. diff --git a/proto/errorpb.proto b/proto/errorpb.proto index f05368e0..97a08d2b 100644 --- a/proto/errorpb.proto +++ b/proto/errorpb.proto @@ -88,8 +88,30 @@ message EpochNotMatch { repeated metapb.Region current_regions = 1; } +// Reasons for rejecting an incompatible transaction protocol request. +enum IncompatibleRequestReason { + IncompatibleRequestReasonUnknown = 0; + // The declared version is outside the compatible range returned in the error. + IncompatibleRequestReasonTxnProtocolVersionOutOfRange = 1; +} + +// A transaction protocol compatibility rejection. +// The compatible range is inclusive: admission reports the store's global range; +// feature checks report [required feature version, store's maximum version]. +message IncompatibleRequest { + IncompatibleRequestReason reason = 1; + // Diagnostic text; clients must use structured fields for control flow. + string message = 2; + // The request's declaration; missing and explicit zero both mean legacy. + uint32 provided_txn_protocol_version = 3; + uint32 min_compatible_txn_protocol_version = 4; + uint32 max_compatible_txn_protocol_version = 5; +} + // ServerIsBusy is the error variant that tells the server is too busy to response. message ServerIsBusy { + // "txn_protocol_incompatible" is a legacy fallback for IncompatibleRequest, + // not an overload signal. string reason = 1; // The suggested backoff time uint64 backoff_ms = 2; @@ -214,4 +236,14 @@ message Error { // UndeterminedResult is the error variant that tells the result is not determined yet. UndeterminedResult undetermined_result = 22; + + // A compatibility rejection with no usable business result. + // For a zero declaration, also set server_is_busy.reason to + // "txn_protocol_incompatible"; omit that fallback for nonzero declarations. + // Clients must handle this error before fallback or ordinary retry logic. + // A client may reselect and resend only when the reason is version out of range, + // the returned range is valid, and provided > max. The selected version must + // satisfy the request's requirements. Such admission rejections must have no + // transaction side effects. All other incompatibilities are terminal. + IncompatibleRequest incompatible_request = 23; } diff --git a/proto/kvrpcpb.proto b/proto/kvrpcpb.proto index d2398e8d..fdea4c28 100644 --- a/proto/kvrpcpb.proto +++ b/proto/kvrpcpb.proto @@ -853,9 +853,16 @@ message MvccGetByStartTsResponse { // Helper messages. +// Client-provided process/product identity for auditing, not authentication or +// authorization. TiDB includes embedded BR; standalone BR uses RequestOriginBR. +// Audit unknown values without mapping them to a known origin or rejecting +// requests solely because of their origin in the initial rollout. enum RequestOrigin { RequestOriginUnknown = 0; RequestOriginTiDB = 1; + RequestOriginTiCDC = 2; + RequestOriginBR = 3; + RequestOriginTiFlash = 4; } // Miscellaneous metadata attached to most requests. @@ -928,11 +935,7 @@ message Context { // Some information used for resource control. ResourceControlContext resource_control_context = 28; - // The declared origin of the request. TiDB server requests set this to TIDB. - // This field is client-provided metadata; consumers should rely on it only when - // the caller is authenticated/authorized by the transport or deployment boundary. - // UNKNOWN means the origin is unset or not recognized by this protocol version. - // Other components should add dedicated enum variants when they need origin-specific behavior. + // Declared process/product origin; see RequestOrigin. RequestOrigin request_origin = 29; // The keyspace that the request is sent to. @@ -966,9 +969,12 @@ message Context { // This field is set by client-go based on an extractor function provided by TiDB. uint64 trace_control_flags = 37; - // TODO(slock): tidb should send requests with a feature flag to indicate whether it - // supports shared locks, so that tikv can fail the requests from old tidb versions - // when needed. + // Transaction protocol version declared for this request; see TxnProtocolVersion. + // Missing or zero means legacy. Any nonzero value promises structured + // IncompatibleRequest handling. BatchCommands declares this per inner request. + // This is a compatibility declaration, not a binary version or authorization. + // Keep uint32 so future version values remain representable. + uint32 txn_protocol_version = 39; } message ResourceControlContext { @@ -1039,6 +1045,20 @@ enum APIVersion { V2 = 2; } +// Transaction protocol version registry. Each version includes all earlier +// semantics. Keep values stable; new versions require joint review by TiKV, +// TiDB and relevant client owners. This registry sets neither client defaults +// nor the store's admission range. +enum TxnProtocolVersion { + // No promise to handle IncompatibleRequest. + TXN_VER_LEGACY = 0; + // Handles IncompatibleRequest according to its error contract. + TXN_VER_SUPPORT_INCOMPATIBLE_ERROR_HANDLING = 1; + // Also supports shared wrapper parsing and expansion, nested lock information + // preservation, and shared-lock pagination and resolve semantics. + TXN_VER_SUPPORT_SHARED_LOCK = 2; +} + message SourceStmt { uint64 start_ts = 1; uint64 connection_id = 2; @@ -1067,9 +1087,8 @@ message LockInfo { // When lock_type is SharedLock, this describes transactions holding the shared lock. // Important: when lock_type is SharedLock, all shared locks must use shared_lock_infos; // DO NOT read from the wrapper LockInfo. - // TODO(slock): tidb should send requests with a feature flag to indicate whether it - // supports shared locks, so that tikv can fail the requests from old tidb versions - // when needed. + // Before returning a shared lock wrapper, the server must check the caller's + // txn_protocol_version and return IncompatibleRequest if support is insufficient. repeated LockInfo shared_lock_infos = 12; // Reserved for file based transaction. diff --git a/proto/metapb.proto b/proto/metapb.proto index 4e38559e..8c7f7984 100644 --- a/proto/metapb.proto +++ b/proto/metapb.proto @@ -67,6 +67,18 @@ message Store { bool physically_destroyed = 12; // NodeState is used to replace StoreState which will be deprecated in the future. NodeState node_state = 13; + // The store's current transaction RPC admission range, not a feature switch or + // a value derived from Store.version. Absence means legacy/unknown; clients + // conservatively use [0, 0] while preserving absence versus an explicit [0, 0]. + // Consumers must fail closed if min > max. TiKV reports this on every startup, + // including upgrades and rollbacks; PD updates and persists it. + TxnProtocolVersionRange txn_protocol_version_range = 14; +} + +// Inclusive transaction protocol version bounds; see kvrpcpb.TxnProtocolVersion. +message TxnProtocolVersionRange { + uint32 min = 1; + uint32 max = 2; } message RegionEpoch { diff --git a/src/common/errors.rs b/src/common/errors.rs index f9a99d2f..85306c24 100644 --- a/src/common/errors.rs +++ b/src/common/errors.rs @@ -14,6 +14,8 @@ use crate::BoundRange; /// future release even if the wire format is compatible. #[doc(inline)] pub use crate::proto::errorpb::Error as ProtoRegionError; +#[doc(inline)] +pub use crate::proto::errorpb::IncompatibleRequest as ProtoIncompatibleRequest; /// Protobuf-generated per-key error returned by TiKV. /// @@ -78,6 +80,9 @@ pub enum Error { /// Errors caused by changes of region information #[error("Region error: {0:?}")] RegionError(Box), + /// The target Store cannot safely interpret this transaction RPC. + #[error("Incompatible transaction request: {0:?}")] + IncompatibleRequest(Box), /// Whether the transaction is committed or not is undetermined #[error("Whether the transaction is committed or not is undetermined")] UndeterminedError(Box), @@ -136,9 +141,54 @@ Use the async TransactionClient instead, or create and use SyncTransactionClient NestedRuntimeError(String), } +/// Priority of direct RPC errors when several responses must be combined. +/// Lower values win; callers retain their own wrapping and cleanup contracts. +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] +pub(crate) enum ErrorPriority { + Undetermined, + Incompatible, + Ordinary, +} + +impl ProtoRegionError { + pub(crate) fn priority(&self) -> ErrorPriority { + if self.undetermined_result.is_some() { + ErrorPriority::Undetermined + } else if self.incompatible_request.is_some() { + ErrorPriority::Incompatible + } else { + ErrorPriority::Ordinary + } + } +} + +impl Error { + pub(crate) fn priority(&self) -> ErrorPriority { + match self { + Self::RegionError(region) => region.priority(), + Self::IncompatibleRequest(_) => ErrorPriority::Incompatible, + _ => ErrorPriority::Ordinary, + } + } +} + +pub(crate) fn is_grpc_error(error: &Error) -> bool { + matches!(error, Error::GrpcAPI(_) | Error::Grpc(_)) +} + +pub(crate) fn is_undetermined_region_error(error: &Error) -> bool { + matches!(error, Error::RegionError(_)) && error.priority() == ErrorPriority::Undetermined +} + impl From for Error { fn from(e: ProtoRegionError) -> Error { - Error::RegionError(Box::new(e)) + if e.priority() == ErrorPriority::Undetermined { + Error::RegionError(Box::new(e)) + } else if let Some(incompatible) = e.incompatible_request { + Error::IncompatibleRequest(Box::new(incompatible)) + } else { + Error::RegionError(Box::new(e)) + } } } diff --git a/src/common/mod.rs b/src/common/mod.rs index 06236c64..2e53053c 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -3,7 +3,11 @@ mod errors; pub mod security; +pub(crate) use self::errors::is_grpc_error; +pub(crate) use self::errors::is_undetermined_region_error; pub use self::errors::Error; +pub(crate) use self::errors::ErrorPriority; +pub use self::errors::ProtoIncompatibleRequest; pub use self::errors::ProtoKeyError; pub use self::errors::ProtoRegionError; pub use self::errors::Result; diff --git a/src/config.rs b/src/config.rs index 92a67e65..3c02a39a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -3,9 +3,34 @@ use std::path::PathBuf; use std::time::Duration; +use crate::{Error, Result}; use serde_derive::Deserialize; use serde_derive::Serialize; +/// Identifies the process issuing transaction RPCs for compatibility auditing. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RequestOrigin { + #[default] + Unknown, + TiDb, + TiCdc, + Br, + TiFlash, +} + +impl RequestOrigin { + pub(crate) fn as_proto(self) -> i32 { + use crate::proto::kvrpcpb::RequestOrigin as ProtoOrigin; + match self { + Self::Unknown => ProtoOrigin::Unknown as i32, + Self::TiDb => ProtoOrigin::TiDb as i32, + Self::TiCdc => ProtoOrigin::TiCdc as i32, + Self::Br => ProtoOrigin::Br as i32, + Self::TiFlash => ProtoOrigin::TiFlash as i32, + } + } +} + /// The configuration for either a [`RawClient`](crate::RawClient) or a /// [`TransactionClient`](crate::TransactionClient). /// @@ -26,10 +51,17 @@ pub struct Config { pub timeout: Duration, pub grpc_max_decoding_message_size: usize, pub keyspace: Option, + #[serde(skip)] + pub(crate) default_txn_protocol_version: u32, + #[serde(skip)] + pub(crate) request_origin: RequestOrigin, } const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 4 * 1024 * 1024; // 4MB +pub(crate) const MAX_SUPPORTED_TXN_PROTOCOL_VERSION: u32 = + crate::proto::kvrpcpb::TxnProtocolVersion::TxnVerSupportIncompatibleErrorHandling as u32; +pub(crate) const DEFAULT_TXN_PROTOCOL_VERSION: u32 = MAX_SUPPORTED_TXN_PROTOCOL_VERSION; impl Default for Config { fn default() -> Self { @@ -40,11 +72,32 @@ impl Default for Config { timeout: DEFAULT_REQUEST_TIMEOUT, grpc_max_decoding_message_size: DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE, keyspace: None, + default_txn_protocol_version: DEFAULT_TXN_PROTOCOL_VERSION, + request_origin: RequestOrigin::Unknown, } } } impl Config { + /// Set the highest transaction protocol version this client can understand. + pub fn with_default_txn_protocol_version(mut self, version: u32) -> Result { + let max = MAX_SUPPORTED_TXN_PROTOCOL_VERSION; + if version > max { + return Err(Error::StringError(format!( + "unsupported transaction protocol version {version}; maximum is {max}" + ))); + } + self.default_txn_protocol_version = version; + Ok(self) + } + + /// Set the process origin attached to transaction RPCs. + #[must_use] + pub fn with_request_origin(mut self, origin: RequestOrigin) -> Self { + self.request_origin = origin; + self + } + /// Set the certificate authority, certificate, and key locations for clients. /// /// By default, this client will use an insecure connection over instead of one protected by @@ -118,3 +171,33 @@ impl Config { self } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transaction_protocol_version_is_bounded() { + assert_eq!(Config::default().default_txn_protocol_version, 1); + assert_eq!( + Config::default() + .with_default_txn_protocol_version(0) + .unwrap() + .default_txn_protocol_version, + 0 + ); + assert!(Config::default() + .with_default_txn_protocol_version(2) + .is_err()); + } + + #[test] + fn compatibility_fields_are_not_deserialized() { + let config: Config = serde_json::from_str( + r#"{"default-txn-protocol-version":0,"request-origin":"TiFlash"}"#, + ) + .unwrap(); + assert_eq!(config.default_txn_protocol_version, 1); + assert_eq!(config.request_origin, RequestOrigin::Unknown); + } +} diff --git a/src/generated/errorpb.rs b/src/generated/errorpb.rs index ddba0b60..eeff78bb 100644 --- a/src/generated/errorpb.rs +++ b/src/generated/errorpb.rs @@ -98,10 +98,31 @@ pub struct EpochNotMatch { #[prost(message, repeated, tag = "1")] pub current_regions: ::prost::alloc::vec::Vec, } +/// A transaction protocol compatibility rejection. +/// The compatible range is inclusive: admission reports the store's global range; +/// feature checks report \[required feature version, store's maximum version\]. +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct IncompatibleRequest { + #[prost(enumeration = "IncompatibleRequestReason", tag = "1")] + pub reason: i32, + /// Diagnostic text; clients must use structured fields for control flow. + #[prost(string, tag = "2")] + pub message: ::prost::alloc::string::String, + /// The request's declaration; missing and explicit zero both mean legacy. + #[prost(uint32, tag = "3")] + pub provided_txn_protocol_version: u32, + #[prost(uint32, tag = "4")] + pub min_compatible_txn_protocol_version: u32, + #[prost(uint32, tag = "5")] + pub max_compatible_txn_protocol_version: u32, +} /// ServerIsBusy is the error variant that tells the server is too busy to response. #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct ServerIsBusy { + /// "txn_protocol_incompatible" is a legacy fallback for IncompatibleRequest, + /// not an overload signal. #[prost(string, tag = "1")] pub reason: ::prost::alloc::string::String, /// The suggested backoff time @@ -274,4 +295,46 @@ pub struct Error { /// UndeterminedResult is the error variant that tells the result is not determined yet. #[prost(message, optional, tag = "22")] pub undetermined_result: ::core::option::Option, + /// A compatibility rejection with no usable business result. + /// For a zero declaration, also set server_is_busy.reason to + /// "txn_protocol_incompatible"; omit that fallback for nonzero declarations. + /// Clients must handle this error before fallback or ordinary retry logic. + /// A client may reselect and resend only when the reason is version out of range, + /// the returned range is valid, and provided > max. The selected version must + /// satisfy the request's requirements. Such admission rejections must have no + /// transaction side effects. All other incompatibilities are terminal. + #[prost(message, optional, tag = "23")] + pub incompatible_request: ::core::option::Option, +} +/// Reasons for rejecting an incompatible transaction protocol request. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum IncompatibleRequestReason { + Unknown = 0, + /// The declared version is outside the compatible range returned in the error. + TxnProtocolVersionOutOfRange = 1, +} +impl IncompatibleRequestReason { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + IncompatibleRequestReason::Unknown => "IncompatibleRequestReasonUnknown", + IncompatibleRequestReason::TxnProtocolVersionOutOfRange => { + "IncompatibleRequestReasonTxnProtocolVersionOutOfRange" + } + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "IncompatibleRequestReasonUnknown" => Some(Self::Unknown), + "IncompatibleRequestReasonTxnProtocolVersionOutOfRange" => { + Some(Self::TxnProtocolVersionOutOfRange) + } + _ => None, + } + } } diff --git a/src/generated/kvrpcpb.rs b/src/generated/kvrpcpb.rs index 89f2d472..209be665 100644 --- a/src/generated/kvrpcpb.rs +++ b/src/generated/kvrpcpb.rs @@ -1290,11 +1290,7 @@ pub struct Context { /// Some information used for resource control. #[prost(message, optional, tag = "28")] pub resource_control_context: ::core::option::Option, - /// The declared origin of the request. TiDB server requests set this to TIDB. - /// This field is client-provided metadata; consumers should rely on it only when - /// the caller is authenticated/authorized by the transport or deployment boundary. - /// UNKNOWN means the origin is unset or not recognized by this protocol version. - /// Other components should add dedicated enum variants when they need origin-specific behavior. + /// Declared process/product origin; see RequestOrigin. #[prost(enumeration = "RequestOrigin", tag = "29")] pub request_origin: i32, /// The keyspace that the request is sent to. @@ -1328,6 +1324,13 @@ pub struct Context { /// This field is set by client-go based on an extractor function provided by TiDB. #[prost(uint64, tag = "37")] pub trace_control_flags: u64, + /// Transaction protocol version declared for this request; see TxnProtocolVersion. + /// Missing or zero means legacy. Any nonzero value promises structured + /// IncompatibleRequest handling. BatchCommands declares this per inner request. + /// This is a compatibility declaration, not a binary version or authorization. + /// Keep uint32 so future version values remain representable. + #[prost(uint32, tag = "39")] + pub txn_protocol_version: u32, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -1390,9 +1393,8 @@ pub struct LockInfo { /// When lock_type is SharedLock, this describes transactions holding the shared lock. /// Important: when lock_type is SharedLock, all shared locks must use shared_lock_infos; /// DO NOT read from the wrapper LockInfo. - /// TODO(slock): tidb should send requests with a feature flag to indicate whether it - /// supports shared locks, so that tikv can fail the requests from old tidb versions - /// when needed. + /// Before returning a shared lock wrapper, the server must check the caller's + /// txn_protocol_version and return IncompatibleRequest if support is insufficient. #[prost(message, repeated, tag = "12")] pub shared_lock_infos: ::prost::alloc::vec::Vec, /// Reserved for file based transaction. @@ -2503,11 +2505,18 @@ impl CommitRole { } } } +/// Client-provided process/product identity for auditing, not authentication or +/// authorization. TiDB includes embedded BR; standalone BR uses RequestOriginBR. +/// Audit unknown values without mapping them to a known origin or rejecting +/// requests solely because of their origin in the initial rollout. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum RequestOrigin { Unknown = 0, TiDb = 1, + TiCdc = 2, + Br = 3, + TiFlash = 4, } impl RequestOrigin { /// String value of the enum field names used in the ProtoBuf definition. @@ -2518,6 +2527,9 @@ impl RequestOrigin { match self { RequestOrigin::Unknown => "RequestOriginUnknown", RequestOrigin::TiDb => "RequestOriginTiDB", + RequestOrigin::TiCdc => "RequestOriginTiCDC", + RequestOrigin::Br => "RequestOriginBR", + RequestOrigin::TiFlash => "RequestOriginTiFlash", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -2525,6 +2537,9 @@ impl RequestOrigin { match value { "RequestOriginUnknown" => Some(Self::Unknown), "RequestOriginTiDB" => Some(Self::TiDb), + "RequestOriginTiCDC" => Some(Self::TiCdc), + "RequestOriginBR" => Some(Self::Br), + "RequestOriginTiFlash" => Some(Self::TiFlash), _ => None, } } @@ -2606,6 +2621,47 @@ impl ApiVersion { } } } +/// Transaction protocol version registry. Each version includes all earlier +/// semantics. Keep values stable; new versions require joint review by TiKV, +/// TiDB and relevant client owners. This registry sets neither client defaults +/// nor the store's admission range. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum TxnProtocolVersion { + /// No promise to handle IncompatibleRequest. + TxnVerLegacy = 0, + /// Handles IncompatibleRequest according to its error contract. + TxnVerSupportIncompatibleErrorHandling = 1, + /// Also supports shared wrapper parsing and expansion, nested lock information + /// preservation, and shared-lock pagination and resolve semantics. + TxnVerSupportSharedLock = 2, +} +impl TxnProtocolVersion { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + TxnProtocolVersion::TxnVerLegacy => "TXN_VER_LEGACY", + TxnProtocolVersion::TxnVerSupportIncompatibleErrorHandling => { + "TXN_VER_SUPPORT_INCOMPATIBLE_ERROR_HANDLING" + } + TxnProtocolVersion::TxnVerSupportSharedLock => "TXN_VER_SUPPORT_SHARED_LOCK", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "TXN_VER_LEGACY" => Some(Self::TxnVerLegacy), + "TXN_VER_SUPPORT_INCOMPATIBLE_ERROR_HANDLING" => { + Some(Self::TxnVerSupportIncompatibleErrorHandling) + } + "TXN_VER_SUPPORT_SHARED_LOCK" => Some(Self::TxnVerSupportSharedLock), + _ => None, + } + } +} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum CommandPri { diff --git a/src/generated/metapb.rs b/src/generated/metapb.rs index 7ee897eb..575c9080 100644 --- a/src/generated/metapb.rs +++ b/src/generated/metapb.rs @@ -57,6 +57,22 @@ pub struct Store { /// NodeState is used to replace StoreState which will be deprecated in the future. #[prost(enumeration = "NodeState", tag = "13")] pub node_state: i32, + /// The store's current transaction RPC admission range, not a feature switch or + /// a value derived from Store.version. Absence means legacy/unknown; clients + /// conservatively use \[0, 0\] while preserving absence versus an explicit \[0, 0\]. + /// Consumers must fail closed if min > max. TiKV reports this on every startup, + /// including upgrades and rollbacks; PD updates and persists it. + #[prost(message, optional, tag = "14")] + pub txn_protocol_version_range: ::core::option::Option, +} +/// Inclusive transaction protocol version bounds; see kvrpcpb.TxnProtocolVersion. +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct TxnProtocolVersionRange { + #[prost(uint32, tag = "1")] + pub min: u32, + #[prost(uint32, tag = "2")] + pub max: u32, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/src/lib.rs b/src/lib.rs index 89b25f45..c185ab8d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -123,6 +123,8 @@ pub use common::security::SecurityManager; #[doc(inline)] pub use common::Error; #[doc(inline)] +pub use common::ProtoIncompatibleRequest; +#[doc(inline)] pub use common::ProtoKeyError; #[doc(inline)] pub use common::ProtoRegionError; @@ -130,6 +132,8 @@ pub use common::ProtoRegionError; pub use common::Result; #[doc(inline)] pub use config::Config; +#[doc(inline)] +pub use config::RequestOrigin; #[doc(inline)] pub use crate::backoff::Backoff; diff --git a/src/mock.rs b/src/mock.rs index 0a1c1b77..e9e07de2 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -6,7 +6,9 @@ //! the system, in particular without requiring a TiKV or PD server, or RPC layer. use std::any::Any; +use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex; use async_trait::async_trait; use derive_new::new; @@ -23,6 +25,7 @@ use crate::region::RegionWithLeader; use crate::store::KvConnect; use crate::store::RegionStore; use crate::store::Request; +use crate::store::TxnProtocolVersionRange; use crate::store::{KvClient, Store}; use crate::Config; use crate::Error; @@ -73,6 +76,75 @@ pub struct MockKvConnect; pub struct MockCluster; +/// Outcomes used to exercise admission retry, rather than just the selector. +#[derive(Clone, Copy, Debug)] +pub enum ReloadOutcome { + Changed, + Unchanged, + Invalid, + Missing, + Failed, +} + +pub fn upper_admission_rejection() -> crate::proto::errorpb::IncompatibleRequest { + crate::proto::errorpb::IncompatibleRequest { + reason: crate::proto::errorpb::IncompatibleRequestReason::TxnProtocolVersionOutOfRange + as i32, + message: "store protocol range changed".into(), + provided_txn_protocol_version: 1, + min_compatible_txn_protocol_version: 0, + max_compatible_txn_protocol_version: 0, + } +} + +/// Models a PD reload updating the metadata read by later routing attempts. +pub struct ProtocolTestClient { + pub pd_client: Arc, + pub reloads: Arc>>, +} + +impl ProtocolTestClient { + pub fn new(client: MockKvClient, outcome: ReloadOutcome) -> Self { + let range = Arc::new(Mutex::new(HashMap::::new())); + let mapped_range = range.clone(); + let mapped_client = client.clone(); + let reloads = Arc::new(Mutex::new(Vec::new())); + let tracked_reloads = reloads.clone(); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_map_region_to_store_hook(move |region| { + let store_id = region.get_store_id()?; + Ok(RegionStore::with_metadata( + region, + store_id, + mapped_range + .lock() + .unwrap() + .get(&store_id) + .copied() + .unwrap_or(TxnProtocolVersionRange { min: 0, max: 1 }), + Arc::new(mapped_client.clone()), + )) + }) + .with_reload_store_hook(move |store_id| { + tracked_reloads.lock().unwrap().push(store_id); + let new_range = match outcome { + ReloadOutcome::Changed => TxnProtocolVersionRange { min: 0, max: 0 }, + ReloadOutcome::Unchanged => TxnProtocolVersionRange { min: 0, max: 1 }, + ReloadOutcome::Invalid => TxnProtocolVersionRange { min: 2, max: 1 }, + ReloadOutcome::Missing => return Ok(None), + ReloadOutcome::Failed => { + return Err(Error::StringError("PD unavailable".into())) + } + }; + range.lock().unwrap().insert(store_id, new_range); + Ok(Some(new_range)) + }), + ); + Self { pd_client, reloads } + } +} + #[allow(clippy::type_complexity)] #[derive(new)] pub struct MockPdClient { @@ -93,6 +165,10 @@ pub struct MockPdClient { /// Optional observer for region cache invalidations. #[new(default)] invalidate_region_hook: Option>, + /// Optional override for conditional Store metadata reloads. + #[new(default)] + reload_store_hook: + Option Result> + Send + Sync + 'static>>, } #[async_trait] @@ -156,6 +232,14 @@ impl MockPdClient { self } + pub fn with_reload_store_hook(mut self, hook: F) -> MockPdClient + where + F: Fn(u64) -> Result> + Send + Sync + 'static, + { + self.reload_store_hook = Some(Arc::new(hook)); + self + } + pub fn region1() -> RegionWithLeader { let mut region = RegionWithLeader::default(); region.region.id = 1; @@ -254,6 +338,16 @@ impl PdClient for MockPdClient { Ok(vec![Store::new(Arc::new(self.client.clone()))]) } + async fn reload_store( + self: Arc, + store_id: u64, + ) -> Result> { + match &self.reload_store_hook { + Some(hook) => hook(store_id), + None => Ok(None), + } + } + async fn get_timestamp(self: Arc) -> Result { Ok(Timestamp::default()) } diff --git a/src/pd/client.rs b/src/pd/client.rs index b8965875..47b9932b 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -2,12 +2,13 @@ use std::collections::HashMap; use std::sync::Arc; +use std::time::{Duration, Instant}; use async_trait::async_trait; use futures::prelude::*; use futures::stream::BoxStream; -use log::info; -use tokio::sync::RwLock; +use log::{info, warn}; +use tokio::sync::{Mutex, OnceCell, RwLock}; use crate::compat::stream_fn; use crate::kv::codec; @@ -25,6 +26,7 @@ use crate::region_cache::RegionCache; use crate::store::KvConnect; use crate::store::RegionStore; use crate::store::TikvConnect; +use crate::store::TxnProtocolVersionRange; use crate::store::{KvClient, Store}; use crate::BoundRange; use crate::Config; @@ -82,6 +84,15 @@ pub trait PdClient: Send + Sync + 'static { async fn all_stores(&self) -> Result>; + /// Conditionally force a Store metadata reload. Implementations coalesce + /// concurrent callers and may suppress attempts during a short cooldown. + async fn reload_store( + self: Arc, + _store_id: StoreId, + ) -> Result> { + Ok(None) + } + fn group_keys_by_region( self: Arc, keys: impl Iterator + Send + Sync + 'static, @@ -217,6 +228,38 @@ pub struct PdRpcClient>>, enable_codec: bool, region_cache: RegionCache>, + store_reloads: Mutex>, +} + +#[derive(Clone)] +struct StoreReload { + started_at: Instant, + result: Arc>, +} + +const STORE_RELOAD_COOLDOWN: Duration = Duration::from_secs(1); + +async fn store_reload_slot( + reloads: &Mutex>, + store_id: StoreId, +) -> StoreReload { + let mut reloads = reloads.lock().await; + match reloads.get(&store_id) { + Some(existing) + if existing.result.get().is_none() + || existing.started_at.elapsed() < STORE_RELOAD_COOLDOWN => + { + existing.clone() + } + _ => { + let reload = StoreReload { + started_at: Instant::now(), + result: Arc::new(OnceCell::new()), + }; + reloads.insert(store_id, reload.clone()); + reload + } + } } #[async_trait] @@ -227,7 +270,12 @@ impl PdClient for PdRpcClient { let store_id = region.get_store_id()?; let store = self.region_cache.get_store_by_id(store_id).await?; let kv_client = self.kv_client(&store.address).await?; - Ok(RegionStore::new(region, Arc::new(kv_client))) + Ok(RegionStore::with_metadata( + region, + store_id, + TxnProtocolVersionRange::from(&store), + Arc::new(kv_client), + )) } async fn region_for_key(&self, key: &Key) -> Result { @@ -252,11 +300,43 @@ impl PdClient for PdRpcClient { let mut stores = Vec::with_capacity(pb_stores.len()); for store in pb_stores { let client = self.kv_client(&store.address).await?; - stores.push(Store::new(Arc::new(client))); + stores.push(Store::with_metadata( + store.id, + TxnProtocolVersionRange::from(&store), + Arc::new(client), + )); } Ok(stores) } + async fn reload_store( + self: Arc, + store_id: StoreId, + ) -> Result> { + let reload = store_reload_slot(&self.store_reloads, store_id).await; + + let refreshed = *reload + .result + .get_or_init(|| async { + match self.region_cache.reload_store_by_id(store_id).await { + Ok(_) => { + warn!("reloaded Store {} after transaction protocol incompatibility", store_id); + true + } + Err(error) => { + warn!("failed to reload Store {} after transaction protocol incompatibility: {:?}", store_id, error); + false + } + } + }) + .await; + if !refreshed { + return Ok(None); + } + let store = self.region_cache.get_store_by_id(store_id).await?; + Ok(Some(TxnProtocolVersionRange::from(&store))) + } + async fn get_timestamp(self: Arc) -> Result { self.pd.clone().get_timestamp().await } @@ -337,6 +417,7 @@ impl PdRpcClient { kv_connect: kv_connect(security_mgr), enable_codec, region_cache: RegionCache::new(pd), + store_reloads: Mutex::new(HashMap::new()), }) } @@ -384,6 +465,26 @@ pub mod test { use crate::store::KvConnect; use crate::Config; + #[tokio::test] + async fn store_reload_slots_are_single_flight_with_attempt_cooldown() { + let reloads = Mutex::new(HashMap::new()); + let first = store_reload_slot(&reloads, 42).await; + let concurrent = store_reload_slot(&reloads, 42).await; + assert!(Arc::ptr_eq(&first.result, &concurrent.result)); + + first.result.set(true).unwrap(); + reloads.lock().await.get_mut(&42).unwrap().started_at = + Instant::now() - STORE_RELOAD_COOLDOWN; + let next_attempt = store_reload_slot(&reloads, 42).await; + assert!(!Arc::ptr_eq(&first.result, &next_attempt.result)); + + let in_flight = store_reload_slot(&reloads, 43).await; + reloads.lock().await.get_mut(&43).unwrap().started_at = + Instant::now() - STORE_RELOAD_COOLDOWN; + let late_waiter = store_reload_slot(&reloads, 43).await; + assert!(Arc::ptr_eq(&in_flight.result, &late_waiter.result)); + } + #[tokio::test] async fn test_kv_client_caching() { let client = block_on(pd_rpc_client()); diff --git a/src/raw/requests.rs b/src/raw/requests.rs index 1f018184..d45e8de8 100644 --- a/src/raw/requests.rs +++ b/src/raw/requests.rs @@ -544,6 +544,20 @@ impl Request for RawCoprocessorRequest { fn set_api_version(&mut self, api_version: kvrpcpb::ApiVersion) { self.inner.set_api_version(api_version); } + + fn txn_protocol_requirement(&self) -> crate::store::TxnProtocolRequirement { + crate::store::TxnProtocolRequirement::NotTransaction + } + + fn prepare_txn_rpc( + &mut self, + range: crate::store::TxnProtocolVersionRange, + default_version: u32, + request_origin: i32, + ) -> Result<()> { + self.inner + .prepare_txn_rpc(range, default_version, request_origin) + } } impl KvRequest for RawCoprocessorRequest { diff --git a/src/region_cache.rs b/src/region_cache.rs index 2a4b0fbb..66e691b6 100644 --- a/src/region_cache.rs +++ b/src/region_cache.rs @@ -166,6 +166,11 @@ impl RegionCache { Ok(store) } + /// Reload Store metadata from PD without invalidating its RPC connection. + pub async fn reload_store_by_id(&self, id: StoreId) -> Result { + self.read_through_store_by_id(id).await + } + pub async fn add_region(&self, region: RegionWithLeader) { // FIXME: will it be the performance bottleneck? let mut cache = self.region_cache.write().await; @@ -306,6 +311,7 @@ mod test { #[derive(Default)] struct MockRetryClient { pub regions: Mutex>, + pub stores: Mutex>, pub get_region_count: AtomicU64, } @@ -343,13 +349,18 @@ mod test { async fn get_store( self: Arc, - _id: crate::region::StoreId, + id: crate::region::StoreId, ) -> Result { - todo!() + self.stores + .lock() + .await + .get(&id) + .cloned() + .ok_or_else(|| Error::StringError(format!("MockRetryClient: store {id} not found"))) } async fn get_all_stores(self: Arc) -> Result> { - todo!() + Ok(self.stores.lock().await.values().cloned().collect()) } async fn get_timestamp(self: Arc) -> Result { @@ -443,6 +454,35 @@ mod test { Ok(()) } + #[tokio::test] + async fn invalid_store_range_remains_visible_on_both_read_paths() -> Result<()> { + let retry_client = Arc::new(MockRetryClient::default()); + let store = metapb::Store { + id: 42, + txn_protocol_version_range: Some(metapb::TxnProtocolVersionRange { min: 2, max: 1 }), + ..Default::default() + }; + retry_client + .stores + .lock() + .await + .insert(store.id, store.clone()); + let cache = RegionCache::new(retry_client); + + let by_id = cache.get_store_by_id(store.id).await?; + let all = cache.read_through_all_stores().await?; + + assert_eq!( + by_id.txn_protocol_version_range, + store.txn_protocol_version_range + ); + assert_eq!( + all[0].txn_protocol_version_range, + store.txn_protocol_version_range + ); + Ok(()) + } + #[tokio::test] async fn test_add_disjoint_regions() { let retry_client = Arc::new(MockRetryClient::default()); diff --git a/src/request/mod.rs b/src/request/mod.rs index 4cac6c50..8584653b 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -23,6 +23,7 @@ pub use self::plan::ResolveLock; pub use self::plan::ResponseWithShard; pub use self::plan::RetryableMultiRegion; pub use self::plan_builder::PlanBuilder; +pub(crate) use self::plan_builder::PlanContext; pub use self::plan_builder::SingleKey; pub use self::shard::Batchable; pub use self::shard::HasNextBatch; @@ -36,6 +37,7 @@ use crate::backoff::PESSIMISTIC_BACKOFF; use crate::store::Request; use crate::store::{HasKeyErrors, Store}; use crate::transaction::HasLocks; +use crate::Result; mod keyspace; pub mod plan; @@ -52,7 +54,11 @@ pub trait KvRequest: Request + Sized + Clone + Sync + Send + 'static { /// For requests or plans which are handled at TiKV store (other than region) level. pub trait StoreRequest { /// Apply the request to specified TiKV store. - fn apply_store(&mut self, store: &Store); + fn apply_store(&mut self, store: &Store) -> Result<()>; + + fn txn_protocol_requirement(&self) -> crate::store::TxnProtocolRequirement { + crate::store::TxnProtocolRequirement::NotTransaction + } } #[derive(Clone, Debug, new, Eq, PartialEq)] @@ -100,6 +106,7 @@ mod test { use super::*; use crate::mock::MockKvClient; use crate::mock::MockPdClient; + use crate::proto::errorpb; use crate::proto::keyspacepb; use crate::proto::kvrpcpb; use crate::proto::metapb::{self, RegionEpoch}; @@ -126,6 +133,10 @@ mod test { } impl HasRegionError for MockRpcResponse { + fn has_region_error(&self) -> bool { + true + } + fn region_error(&mut self) -> Option { Some(crate::proto::errorpb::Error::default()) } @@ -157,6 +168,19 @@ mod test { } fn set_api_version(&mut self, _: kvrpcpb::ApiVersion) {} + + fn txn_protocol_requirement(&self) -> crate::store::TxnProtocolRequirement { + crate::store::TxnProtocolRequirement::NotTransaction + } + + fn prepare_txn_rpc( + &mut self, + _: crate::store::TxnProtocolVersionRange, + _: u32, + _: i32, + ) -> Result<()> { + Ok(()) + } } #[async_trait] @@ -222,6 +246,10 @@ mod test { } impl HasRegionError for MockOkResponse { + fn has_region_error(&self) -> bool { + false + } + fn region_error(&mut self) -> Option { None } @@ -357,6 +385,19 @@ mod test { } fn set_api_version(&mut self, _: kvrpcpb::ApiVersion) {} + + fn txn_protocol_requirement(&self) -> crate::store::TxnProtocolRequirement { + crate::store::TxnProtocolRequirement::NotTransaction + } + + fn prepare_txn_rpc( + &mut self, + _: crate::store::TxnProtocolVersionRange, + _: u32, + _: i32, + ) -> Result<()> { + Ok(()) + } } #[async_trait] @@ -530,6 +571,10 @@ mod test { } impl HasRegionError for MockOkResponse { + fn has_region_error(&self) -> bool { + false + } + fn region_error(&mut self) -> Option { None } @@ -638,6 +683,19 @@ mod test { } fn set_api_version(&mut self, _: kvrpcpb::ApiVersion) {} + + fn txn_protocol_requirement(&self) -> crate::store::TxnProtocolRequirement { + crate::store::TxnProtocolRequirement::NotTransaction + } + + fn prepare_txn_rpc( + &mut self, + _: crate::store::TxnProtocolVersionRange, + _: u32, + _: i32, + ) -> Result<()> { + Ok(()) + } } #[async_trait] @@ -805,6 +863,176 @@ mod test { .expect("request context") } + #[rstest::rstest] + #[case(false, false)] + #[case(true, false)] + #[case(false, true)] + #[case(true, true)] + #[tokio::test] + async fn fallback_preserves_high_priority_errors( + #[case] not_leader: bool, + #[case] undetermined: bool, + ) { + let incompatible = errorpb::IncompatibleRequest { + reason: errorpb::IncompatibleRequestReason::TxnProtocolVersionOutOfRange as i32, + message: "legacy declaration rejected".into(), + provided_txn_protocol_version: 0, + min_compatible_txn_protocol_version: 1, + max_compatible_txn_protocol_version: 2, + }; + let mixed = errorpb::Error { + not_leader: not_leader.then_some(errorpb::NotLeader::default()), + server_is_busy: (!not_leader).then_some(errorpb::ServerIsBusy { + reason: "txn_protocol_incompatible".into(), + ..Default::default() + }), + undetermined_result: undetermined.then_some(errorpb::UndeterminedResult { + message: "apply outcome unknown".into(), + }), + incompatible_request: (!undetermined).then_some(incompatible.clone()), + ..Default::default() + }; + let expected = mixed.clone(); + let accesses = Arc::new(std::sync::Mutex::new(Vec::new())); + let dispatch_accesses = accesses.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + let context = fallback_get_context(request); + assert_eq!(context.txn_protocol_version, 0); + let store_id = context.peer.as_ref().unwrap().store_id; + dispatch_accesses.lock().unwrap().push(store_id); + match store_id { + 41 => Err(Error::GrpcAPI(tonic::Status::deadline_exceeded( + "cached leader", + ))), + 44 => Ok(Box::new(kvrpcpb::GetResponse { + region_error: Some(mixed.clone()), + ..Default::default() + }) as Box), + 45 => Ok(Box::new(kvrpcpb::GetResponse::default()) as Box), + _ => panic!("unexpected voter"), + } + }); + let invalidations = Arc::new(AtomicUsize::new(0)); + let invalidated = invalidations.clone(); + let leader_updates = Arc::new(AtomicUsize::new(0)); + let updated = leader_updates.clone(); + let reloads = Arc::new(AtomicUsize::new(0)); + let reloaded = reloads.clone(); + let region = fallback_region(&[44, 45]); + let pd_client = Arc::new( + MockPdClient::new(client) + .with_region_for_key_hook(move |_| Ok(region.clone())) + .with_invalidate_region_hook(move |_| { + invalidated.fetch_add(1, Ordering::SeqCst); + }) + .with_update_leader_hook(move |_, _| { + updated.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + .with_reload_store_hook(move |_| { + reloaded.fetch_add(1, Ordering::SeqCst); + Ok(None) + }), + ); + let result = PlanBuilder::new( + pd_client, + Keyspace::Disable, + kvrpcpb::GetRequest { + key: b"key".to_vec(), + ..Default::default() + }, + ) + .retry_multi_region_terminal_on_undetermined(Backoff::no_backoff()) + .plan() + .execute() + .await; + // The old predicates continue to voter 45 and mask the error with success. + assert_eq!(*accesses.lock().unwrap(), vec![41, 44]); + match result { + Err(Error::RegionError(error)) if undetermined => assert_eq!(*error, expected), + Err(Error::IncompatibleRequest(error)) if !undetermined => { + assert_eq!(*error, incompatible); + } + other => panic!("wrong outcome: {other:?}"), + } + assert_eq!(invalidations.load(Ordering::SeqCst), 0); + assert_eq!(leader_updates.load(Ordering::SeqCst), 0); + assert_eq!(reloads.load(Ordering::SeqCst), 0); + } + + #[rstest::rstest] + #[case(crate::mock::ReloadOutcome::Changed)] + #[case(crate::mock::ReloadOutcome::Unchanged)] + #[tokio::test] + async fn admission_reload_retries_only_the_rejected_shard( + #[case] outcome: crate::mock::ReloadOutcome, + ) { + let sent = Arc::new(std::sync::Mutex::new(Vec::new())); + let recorded = sent.clone(); + let incompatible = crate::mock::upper_admission_rejection(); + let rejection = incompatible.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + let request = request.downcast_ref::().unwrap(); + let context = request.context.as_ref().unwrap(); + let store_id = context.peer.as_ref().unwrap().store_id; + recorded + .lock() + .unwrap() + .push((store_id, context.txn_protocol_version)); + Ok(Box::new(kvrpcpb::BatchGetResponse { + pairs: if store_id == 41 && context.txn_protocol_version == 1 { + vec![kvrpcpb::KvPair { + error: Some(kvrpcpb::KeyError { + locked: Some(kvrpcpb::LockInfo { + key: vec![1], + primary_lock: vec![1], + lock_version: 1, + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }] + } else { + vec![] + }, + region_error: (store_id == 41 && context.txn_protocol_version == 1).then_some( + errorpb::Error { + incompatible_request: Some(rejection.clone()), + ..Default::default() + }, + ), + ..Default::default() + }) as Box) + }); + let fixture = crate::mock::ProtocolTestClient::new(client, outcome); + let result = PlanBuilder::new( + fixture.pd_client, + Keyspace::Disable, + kvrpcpb::BatchGetRequest { + keys: vec![vec![1], vec![11]], + ..Default::default() + }, + ) + .resolve_lock(Timestamp::default(), Backoff::no_jitter_backoff(0, 0, 1)) + .retry_multi_region(Backoff::no_jitter_backoff(0, 0, 1)) + .plan() + .execute() + .await; + let mut versions = sent.lock().unwrap().clone(); + versions.sort_unstable(); + if matches!(outcome, crate::mock::ReloadOutcome::Changed) { + assert!(result.is_ok()); + assert_eq!(versions, vec![(41, 0), (41, 1), (42, 1)]); + } else { + assert!( + matches!(result, Err(Error::IncompatibleRequest(error)) if *error == incompatible) + ); + assert_eq!(versions, vec![(41, 1), (42, 1)]); + } + assert_eq!(*fixture.reloads.lock().unwrap(), vec![41]); + } + async fn execute_fallback_get( pd_client: Arc, backoff: Backoff, diff --git a/src/request/plan.rs b/src/request/plan.rs index e1a1eb5b..a2c5bd58 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -16,6 +16,8 @@ use tokio::time::sleep; use tonic::Code; use crate::backoff::Backoff; +use crate::common::is_grpc_error; +use crate::common::ErrorPriority; use crate::pd::PdClient; use crate::proto::errorpb; use crate::proto::errorpb::EpochNotMatch; @@ -29,12 +31,13 @@ use crate::request::NextBatch; use crate::request::Shardable; use crate::request::{KvRequest, StoreRequest}; use crate::stats::tikv_stats; +use crate::store::reload_txn_protocol_range; use crate::store::HasRegionError; use crate::store::HasRegionErrors; use crate::store::KvClient; use crate::store::RegionStore; +use crate::store::TxnProtocolRequirement; use crate::store::{HasKeyErrors, Store}; -use crate::transaction::resolve_locks; use crate::transaction::HasLocks; use crate::transaction::ResolveLocksContext; use crate::transaction::ResolveLocksOptions; @@ -42,7 +45,9 @@ use crate::util::iter::FlatMapOkIterExt; use crate::Error; use crate::Result; +#[cfg(test)] use super::keyspace::Keyspace; +use super::plan_builder::PlanContext; /// A plan for how to execute a request. A user builds up a plan with various /// options, then exectutes it. @@ -60,6 +65,7 @@ pub trait Plan: Sized + Clone + Sync + Send + 'static { pub struct Dispatch { pub request: Req, pub kv_client: Option>, + pub(crate) context: PlanContext, } #[async_trait] @@ -83,19 +89,24 @@ impl Plan for Dispatch { } impl StoreRequest for Dispatch { - fn apply_store(&mut self, store: &Store) { + fn apply_store(&mut self, store: &Store) -> Result<()> { self.kv_client = Some(store.client.clone()); - self.request.apply_store(store); + self.request.prepare_txn_rpc( + store.txn_protocol_version_range, + self.context.default_txn_protocol_version(), + self.context.request_origin().as_proto(), + )?; + self.request.apply_store(store) + } + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + crate::store::Request::txn_protocol_requirement(&self.request) } } const MULTI_REGION_CONCURRENCY: usize = 16; const MULTI_STORES_CONCURRENCY: usize = 16; -pub(crate) fn is_grpc_error(e: &Error) -> bool { - matches!(e, Error::GrpcAPI(_) | Error::Grpc(_)) -} - /// Await every task in `join_set`, reassembling the results in spawn order. /// /// Contract: on any join failure (a panicked or cancelled task) the remaining tasks @@ -134,17 +145,10 @@ where .collect()) } -/// Did the server say the request's outcome is UNKNOWN — e.g. a raft timeout where the -/// apply result was never observed (`errorpb.UndeterminedResult`)? A commit receiving -/// this must be reported as undetermined: reporting plain failure would invite the -/// caller to retry effects that may already be durable. -pub(crate) fn is_undetermined_region_error(e: &Error) -> bool { - matches!(e, Error::RegionError(re) if re.undetermined_result.is_some()) -} - pub struct RetryableMultiRegion { pub(super) inner: P, pub pd_client: Arc, + pub(crate) context: PlanContext, pub backoff: Backoff, /// Preserve all regions' results for other downstream plans to handle. @@ -201,8 +205,15 @@ where } } + fn has_high_priority_region_error(&self) -> bool { + self.region_error + .as_ref() + .is_some_and(|error| error.priority() != ErrorPriority::Ordinary) + } + fn is_fallback_not_leader(&self) -> bool { - self.used_fallback + !self.has_high_priority_region_error() + && self.used_fallback && self.key_errors.is_none() && self .region_error @@ -211,7 +222,8 @@ where } fn is_fallback_server_busy(&self) -> bool { - self.used_fallback + !self.has_high_priority_region_error() + && self.used_fallback && self.key_errors.is_none() && self .region_error @@ -270,6 +282,7 @@ impl CandidateState { struct RetryContext { pd_client: Arc, + context: PlanContext, permits: Arc, preserve_region_results: bool, terminal_on_undetermined: bool, @@ -280,6 +293,7 @@ impl Clone for RetryContext { fn clone(&self) -> Self { Self { pd_client: self.pd_client.clone(), + context: self.context, permits: self.permits.clone(), preserve_region_results: self.preserve_region_results, terminal_on_undetermined: self.terminal_on_undetermined, @@ -360,12 +374,12 @@ where // A plan may involve multiple shards #[async_recursion] async fn single_plan_handler( - context: RetryContext, + retry_context: RetryContext, current_plan: P, backoff: Backoff, ) -> Result<::Result> { let shards = current_plan - .shards(&context.pd_client) + .shards(&retry_context.pd_client) .collect::>() .await; let shards_len = shards.len(); @@ -380,13 +394,13 @@ where } }; let clone = current_plan.clone_then_apply_shard(shard); - let shard_context = context.clone(); + let shard_retry_context = retry_context.clone(); let backoff = backoff.clone(); join_set.spawn(async move { ( idx, Self::single_shard_handler( - shard_context, + shard_retry_context, clone, region, backoff, @@ -399,7 +413,7 @@ where let results = collect_join_set_results(join_set, shards_len, "single_plan_handler").await?; - if context.preserve_region_results { + if retry_context.preserve_region_results { Ok(results .into_iter() .flat_map_ok(|x| x) @@ -417,19 +431,21 @@ where // that may already be durable. So an undetermined error // from ANY shard wins over a determinate one. let mut oks = Vec::with_capacity(results.len()); - let mut first_err: Option = None; - let mut undetermined: Option = None; + let mut selected_error: Option = None; for r in results { match r { Ok(v) => oks.push(v), - Err(e) if undetermined.is_none() && is_undetermined_region_error(&e) => { - undetermined = Some(e) + Err(error) + if selected_error + .as_ref() + .is_none_or(|selected| error.priority() < selected.priority()) => + { + selected_error = Some(error) } - Err(e) if first_err.is_none() => first_err = Some(e), Err(_) => {} } } - if let Some(e) = undetermined.or(first_err) { + if let Some(e) = selected_error { return Err(e); } Ok(oks.into_iter().flatten().collect()) @@ -449,7 +465,7 @@ where let mut candidate_region = region.clone(); candidate_region.leader = Some(peer.clone()); - let region_store = match pd_client + let mut region_store = match pd_client .clone() .map_region_to_store(candidate_region) .await @@ -467,6 +483,16 @@ where return CandidateRoundResult::MapRegionToStoreError(err); } }; + if let TxnProtocolRequirement::Transaction { required_version } = + plan.txn_protocol_requirement() + { + if required_version > region_store.txn_protocol_version_range.max { + if let Ok(Some(range)) = pd_client.clone().reload_store(region_store.store_id).await + { + region_store.txn_protocol_version_range = range; + } + } + } if let Err(error) = plan.apply_store(®ion_store) { // Applying a successfully mapped Store mutates the request. This is // not a PD/connection lookup failure and trying another peer cannot @@ -646,7 +672,7 @@ where } async fn retry_same_region( - context: RetryContext, + retry_context: RetryContext, plan: P, region: RegionWithLeader, mut backoff: Backoff, @@ -656,7 +682,8 @@ where match backoff.next_delay_duration() { Some(duration) => { sleep(duration).await; - Self::single_shard_handler(context, plan, region, backoff, candidate_state).await + Self::single_shard_handler(retry_context, plan, region, backoff, candidate_state) + .await } None => Err(error), } @@ -664,7 +691,7 @@ where #[async_recursion] async fn single_shard_handler( - context: RetryContext, + retry_context: RetryContext, mut plan: P, region: RegionWithLeader, backoff: Backoff, @@ -679,11 +706,11 @@ where let followers_only = candidate_state.take_follower_probe(); let response = match Self::execute_candidate_round( - &context.pd_client, + &retry_context.pd_client, &mut plan, ®ion, - &context.permits, - context.terminal_on_dispatch_error, + &retry_context.permits, + retry_context.terminal_on_dispatch_error, followers_only, ) .await @@ -691,7 +718,7 @@ where CandidateRoundResult::Response(response) => *response, CandidateRoundResult::MapRegionToStoreError(error) if followers_only => { return Self::retry_same_region( - context, + retry_context, plan, region, backoff, @@ -705,7 +732,7 @@ where // The cached peer list may have been replaced completely, so // reload the Region instead of retrying the same stale peers. return Self::retry_after_routing_error( - context, + retry_context, plan, region_ver_id, backoff, @@ -722,12 +749,18 @@ where // and restore the cached leader without charging another // backoff: the one-shot probe cannot fire again in this // request, so the next leader error follows the normal path. - return Self::single_shard_handler(context, plan, region, backoff, candidate_state) - .await; + return Self::single_shard_handler( + retry_context, + plan, + region, + backoff, + candidate_state, + ) + .await; } CandidateRoundResult::RoutingError { error, .. } if followers_only => { return Self::retry_same_region( - context, + retry_context, plan, region, backoff, @@ -740,14 +773,14 @@ where error, invalidate_region: false, } => { - return Self::retry_after_error(context, plan, backoff, error).await; + return Self::retry_after_error(retry_context, plan, backoff, error).await; } CandidateRoundResult::RoutingError { error, invalidate_region: true, } => { return Self::retry_after_routing_error( - context, + retry_context, plan, region_ver_id, backoff, @@ -758,7 +791,7 @@ where }; Self::handle_candidate_response( - context, + retry_context, plan, region, backoff, @@ -770,7 +803,7 @@ where } async fn handle_candidate_response( - context: RetryContext, + retry_context: RetryContext, plan: P, region: RegionWithLeader, backoff: Backoff, @@ -780,8 +813,12 @@ where ) -> Result<::Result> { let region_ver_id = region.ver_id(); let fallback_leader = response.accepted_fallback_leader(); - Self::update_leader_after_fallback(&context.pd_client, ®ion_ver_id, fallback_leader) - .await; + Self::update_leader_after_fallback( + &retry_context.pd_client, + ®ion_ver_id, + fallback_leader, + ) + .await; let CandidateResponse { response, key_errors, @@ -790,6 +827,48 @@ where used_fallback, } = response; + // Region-level uncertainty and incompatibility outrank key-level and + // load-shedding fields when an older/newer server populates more than + // one field in the same response. + if let Some(error) = region_error.as_ref() { + if error.undetermined_result.is_some() { + return Self::handle_region_response( + retry_context, + plan, + region.ver_id(), + region_store, + backoff, + error.clone(), + ) + .await; + } + if let Some(incompatible) = error.incompatible_request.as_ref() { + let mut backoff = backoff; + if reload_txn_protocol_range( + retry_context.pd_client.clone(), + region_store.store_id, + region_store.txn_protocol_version_range, + retry_context.context.default_txn_protocol_version(), + plan.txn_protocol_requirement(), + incompatible, + &mut backoff, + ) + .await + .is_some() + { + return Self::single_shard_handler( + retry_context, + plan, + region, + backoff, + candidate_state, + ) + .await; + } + return Err(Error::IncompatibleRequest(Box::new(incompatible.clone()))); + } + } + if let Some(error) = key_errors { debug!("single_shard_handler:execute: key errors: {:?}", error); return Ok(vec![Err(Error::MultipleKeyErrors(error))]); @@ -805,7 +884,7 @@ where candidate_state.record_leader_busy(peer.id, server_is_busy.estimated_wait_ms); } return Self::retry_same_region( - context, + retry_context, plan, region, backoff, @@ -819,7 +898,7 @@ where // The one-shot probe was inconclusive. Restore the cached // leader and resume the normal ServerIsBusy backoff path. return Self::retry_same_region( - context, + retry_context, plan, region, backoff, @@ -830,12 +909,19 @@ where } } - Self::handle_region_response(context, plan, region.ver_id(), region_store, backoff, error) - .await + Self::handle_region_response( + retry_context, + plan, + region.ver_id(), + region_store, + backoff, + error, + ) + .await } async fn handle_region_response( - context: RetryContext, + retry_context: RetryContext, plan: P, region_ver_id: RegionVerId, region_store: RegionStore, @@ -849,18 +935,25 @@ where // For CAS and commit points, an unknown apply outcome must surface on // first sight: replaying could contradict its own effect, and a later // different error must not overwrite the uncertainty. - if context.terminal_on_undetermined && error.undetermined_result.is_some() { + if retry_context.terminal_on_undetermined && error.undetermined_result.is_some() { return Err(Error::RegionError(Box::new(error))); } + if error.undetermined_result.is_none() { + if let Some(incompatible) = error.incompatible_request.as_ref() { + return Err(Error::IncompatibleRequest(Box::new(incompatible.clone()))); + } + } + match backoff.next_delay_duration() { Some(duration) => { let region_error_resolved = - handle_region_error(context.pd_client.clone(), error, region_store).await?; + handle_region_error(retry_context.pd_client.clone(), error, region_store) + .await?; if !region_error_resolved { sleep(duration).await; } - Self::single_plan_handler(context, plan, backoff).await + Self::single_plan_handler(retry_context, plan, backoff).await } None => { warn!( @@ -873,19 +966,22 @@ where } async fn retry_after_routing_error( - context: RetryContext, + retry_context: RetryContext, plan: P, region: RegionVerId, backoff: Backoff, error: Error, ) -> Result<::Result> { debug!("retry_after_routing_error: {:?}", error); - context.pd_client.invalidate_region_cache(region).await; - Self::retry_after_error(context, plan, backoff, error).await + retry_context + .pd_client + .invalidate_region_cache(region) + .await; + Self::retry_after_error(retry_context, plan, backoff, error).await } async fn retry_after_error( - context: RetryContext, + retry_context: RetryContext, plan: P, mut backoff: Backoff, error: Error, @@ -893,7 +989,7 @@ where match backoff.next_delay_duration() { Some(duration) => { sleep(duration).await; - Self::single_plan_handler(context, plan, backoff).await + Self::single_plan_handler(retry_context, plan, backoff).await } None => Err(error), } @@ -912,6 +1008,14 @@ pub(crate) async fn handle_region_error( let ver_id = region_store.region_with_leader.ver_id(); let store_id = region_store.region_with_leader.get_store_id(); debug!("handling region error: {:?}, region: {:?}", e, ver_id); + if e.undetermined_result.is_some() { + // Preserve the highest-priority ambiguous outcome even if a malformed + // response also populated another region-error field. + return Ok(false); + } + if let Some(incompatible) = e.incompatible_request.as_ref() { + return Err(Error::IncompatibleRequest(Box::new(incompatible.clone()))); + } if let Some(not_leader) = e.not_leader { if let Some(leader) = not_leader.leader { match pd_client @@ -943,15 +1047,6 @@ pub(crate) async fn handle_region_error( } else if e.stale_command.is_some() || e.region_not_found.is_some() { pd_client.invalidate_region_cache(ver_id).await; Ok(false) - } else if e.undetermined_result.is_some() { - // The apply outcome is UNKNOWN (a raft timeout, errorpb.UndeterminedResult). - // Default: retry — matching client-go's ACTION layers (ordinary prewrites and - // secondary commits back off and re-send; re-applying an idempotent request - // resolves the uncertainty). Routing is not suspect, so nothing is - // invalidated. On backoff exhaustion the error escapes UNCHANGED, and commit - // paths classify it via `is_undetermined_region_error`. Plans for which a - // replay is unsafe never reach this arm — see `terminal_on_undetermined`. - Ok(false) } else if e.server_is_busy.is_some() { // ServerIsBusy is a definitive rejection, so retrying is safe. The // cached leader remains valid; the caller applies the normal Region @@ -1017,6 +1112,7 @@ impl Clone for RetryableMultiRegion { RetryableMultiRegion { inner: self.inner.clone(), pd_client: self.pd_client.clone(), + context: self.context, backoff: self.backoff.clone(), preserve_region_results: self.preserve_region_results, terminal_on_undetermined: self.terminal_on_undetermined, @@ -1037,20 +1133,22 @@ where // too many concurrent requests, TiKV is more likely to return a "TiKV // is busy" error let concurrency_permits = Arc::new(Semaphore::new(MULTI_REGION_CONCURRENCY)); - let context = RetryContext { + let retry_context = RetryContext { pd_client: self.pd_client.clone(), + context: self.context, permits: concurrency_permits, preserve_region_results: self.preserve_region_results, terminal_on_undetermined: self.terminal_on_undetermined, terminal_on_dispatch_error: self.terminal_on_dispatch_error, }; - Self::single_plan_handler(context, self.inner.clone(), self.backoff.clone()).await + Self::single_plan_handler(retry_context, self.inner.clone(), self.backoff.clone()).await } } pub struct RetryableAllStores { pub(super) inner: P, pub pd_client: Arc, + pub(crate) context: PlanContext, pub backoff: Backoff, } @@ -1059,6 +1157,7 @@ impl Clone for RetryableAllStores { RetryableAllStores { inner: self.inner.clone(), pd_client: self.pd_client.clone(), + context: self.context, backoff: self.backoff.clone(), } } @@ -1080,15 +1179,39 @@ where let stores = self.pd_client.clone().all_stores().await?; let stores_len = stores.len(); let mut join_set = JoinSet::new(); - for (idx, store) in stores.into_iter().enumerate() { + for (idx, mut store) in stores.into_iter().enumerate() { let mut clone = self.inner.clone(); - clone.apply_store(&store); + if let TxnProtocolRequirement::Transaction { required_version } = + clone.txn_protocol_requirement() + { + if required_version > store.txn_protocol_version_range.max { + if let Ok(Some(range)) = + self.pd_client.clone().reload_store(store.store_id).await + { + store.txn_protocol_version_range = range; + } + } + } + if let Err(error) = clone.apply_store(&store) { + join_set.spawn(async move { (idx, Err(error)) }); + continue; + } let backoff = self.backoff.clone(); let concurrency_permits = concurrency_permits.clone(); + let pd_client = self.pd_client.clone(); + let context = self.context; join_set.spawn(async move { ( idx, - Self::single_store_handler(clone, backoff, concurrency_permits).await, + Self::single_store_handler( + clone, + store, + pd_client, + context, + backoff, + concurrency_permits, + ) + .await, ) }); } @@ -1099,12 +1222,15 @@ where } } -impl RetryableAllStores +impl RetryableAllStores where P::Result: HasKeyErrors + HasRegionError, { async fn single_store_handler( - plan: P, + mut plan: P, + mut store: Store, + pd_client: Arc, + context: PlanContext, mut backoff: Backoff, permits: Arc, ) -> Result { @@ -1115,9 +1241,33 @@ where match res { Ok(mut resp) => { + let region_error = resp.region_error(); + if let Some(e) = region_error.as_ref() { + if e.undetermined_result.is_some() { + return Err(Error::RegionError(Box::new(e.clone()))); + } + if let Some(incompatible) = e.incompatible_request.as_ref() { + if let Some(range) = reload_txn_protocol_range( + pd_client.clone(), + store.store_id, + store.txn_protocol_version_range, + context.default_txn_protocol_version(), + plan.txn_protocol_requirement(), + incompatible, + &mut backoff, + ) + .await + { + store.txn_protocol_version_range = range; + plan.apply_store(&store)?; + continue; + } + return Err(Error::IncompatibleRequest(Box::new(incompatible.clone()))); + } + } if let Some(e) = resp.key_errors() { return Err(Error::MultipleKeyErrors(e)); - } else if let Some(e) = resp.region_error() { + } else if let Some(e) = region_error { // Store request should not return region error. return Err(Error::RegionError(Box::new(e))); } else { @@ -1235,7 +1385,7 @@ pub struct ResolveLock { pub timestamp: Timestamp, pub pd_client: Arc, pub backoff: Backoff, - pub keyspace: Keyspace, + pub(crate) context: PlanContext, } impl Clone for ResolveLock { @@ -1245,7 +1395,7 @@ impl Clone for ResolveLock { timestamp: self.timestamp.clone(), pd_client: self.pd_client.clone(), backoff: self.backoff.clone(), - keyspace: self.keyspace, + context: self.context, } } } @@ -1253,7 +1403,7 @@ impl Clone for ResolveLock { #[async_trait] impl Plan for ResolveLock where - P::Result: HasLocks, + P::Result: HasLocks + HasRegionError, { type Result = P::Result; @@ -1261,6 +1411,13 @@ where let mut result = self.inner.execute().await?; let mut clone = self.clone(); loop { + // Region errors belong to the outer retry/error-handling layer. + // Preserve the entire response before consuming locks or issuing + // auxiliary RPCs, including after a lock-resolution retry. + if result.has_region_error() { + return Ok(result); + } + let locks = result.take_locks(); if locks.is_empty() { return Ok(result); @@ -1271,11 +1428,11 @@ where } let pd_client = self.pd_client.clone(); - let live_locks = resolve_locks( + let live_locks = crate::transaction::resolve_locks_with_context( locks, self.timestamp.clone(), pd_client.clone(), - self.keyspace, + self.context, ) .await?; if live_locks.is_empty() { @@ -1310,6 +1467,10 @@ impl Clone for CleanupLocksResult { } impl HasRegionError for CleanupLocksResult { + fn has_region_error(&self) -> bool { + self.region_error.is_some() + } + fn region_error(&mut self) -> Option { self.region_error.take() } @@ -1342,8 +1503,8 @@ pub struct CleanupLocks { pub options: ResolveLocksOptions, pub store: Option, pub pd_client: Arc, - pub keyspace: Keyspace, pub backoff: Backoff, + pub(crate) context: PlanContext, } impl Clone for CleanupLocks { @@ -1354,8 +1515,8 @@ impl Clone for CleanupLocks { options: self.options, store: None, pd_client: self.pd_client.clone(), - keyspace: self.keyspace, backoff: self.backoff.clone(), + context: self.context, } } } @@ -1370,7 +1531,8 @@ where async fn execute(&self) -> Result { let mut result = CleanupLocksResult::default(); let mut inner = self.inner.clone(); - let mut lock_resolver = crate::transaction::LockResolver::new(self.ctx.clone()); + let mut lock_resolver = + crate::transaction::LockResolver::new_with_context(self.ctx.clone(), self.context); let region = &self.store.as_ref().unwrap().region_with_leader; let mut has_more_batch = true; @@ -1378,14 +1540,14 @@ where let mut scan_lock_resp = inner.execute().await?; // Propagate errors to `retry_multi_region` for retry. - if let Some(e) = scan_lock_resp.key_errors() { - info!("CleanupLocks::execute, inner key errors:{:?}", e); - result.key_error = Some(e); - return Ok(result); - } else if let Some(e) = scan_lock_resp.region_error() { + if let Some(e) = scan_lock_resp.region_error() { info!("CleanupLocks::execute, inner region error:{}", e.message); result.region_error = Some(e); return Ok(result); + } else if let Some(e) = scan_lock_resp.key_errors() { + info!("CleanupLocks::execute, inner key errors:{:?}", e); + result.key_error = Some(e); + return Ok(result); } // Iterate to next batch of inner. @@ -1423,7 +1585,7 @@ where self.store.clone().unwrap(), locks, self.pd_client.clone(), - self.keyspace, + self.context.keyspace(), ) .await { @@ -1432,10 +1594,13 @@ where } Err(Error::ExtractedErrors(mut errors)) => { // Propagate errors to `retry_multi_region` for retry. - if let Error::RegionError(e) = errors.pop().unwrap() { - result.region_error = Some(*e); - } else { - result.key_error = Some(errors); + match errors.pop() { + Some(Error::RegionError(e)) => result.region_error = Some(*e), + Some(error) => { + errors.push(error); + result.key_error = Some(errors); + } + None => return Err(Error::ExtractedErrors(errors)), } return Ok(result); } @@ -1483,15 +1648,18 @@ where async fn execute(&self) -> Result { let mut result = self.inner.execute().await?; - if let Some(errors) = result.key_errors() { + if let Some(errors) = result.region_errors() { + let mut errors = errors.into_iter().map(Error::from).collect::>(); + errors.sort_by_key(Error::priority); + if errors + .first() + .is_some_and(|error| error.priority() != ErrorPriority::Ordinary) + { + return Err(errors.remove(0)); + } + Err(Error::ExtractedErrors(errors)) + } else if let Some(errors) = result.key_errors() { Err(Error::ExtractedErrors(errors)) - } else if let Some(errors) = result.region_errors() { - Err(Error::ExtractedErrors( - errors - .into_iter() - .map(|e| Error::RegionError(Box::new(e))) - .collect(), - )) } else { Ok(result) } @@ -1552,6 +1720,10 @@ impl HasLocks for ResponseWithShard { } impl HasRegionError for ResponseWithShard { + fn has_region_error(&self) -> bool { + self.0.has_region_error() + } + fn region_error(&mut self) -> Option { self.0.region_error() } @@ -1567,9 +1739,296 @@ mod test { use futures::stream::{self}; use super::*; + use crate::common::is_undetermined_region_error; use crate::mock::MockPdClient; use crate::proto::kvrpcpb::BatchGetResponse; + #[derive(Clone)] + struct RegionErrorsPlan(Vec); + + #[derive(Clone)] + struct LockResponsePlan { + responses: Vec, + attempts: Arc, + } + + #[async_trait] + impl Plan for LockResponsePlan { + type Result = kvrpcpb::GetResponse; + + async fn execute(&self) -> Result { + let attempt = self.attempts.fetch_add(1, Ordering::SeqCst); + Ok(self + .responses + .get(attempt) + .expect("unexpected request replay") + .clone()) + } + } + + fn mixed_lock_response(priority: ErrorPriority) -> kvrpcpb::GetResponse { + let mut region_error = errorpb::Error { + message: "original region error".into(), + ..Default::default() + }; + match priority { + ErrorPriority::Undetermined => { + region_error.undetermined_result = Some(Default::default()) + } + ErrorPriority::Incompatible => { + region_error.incompatible_request = Some(Default::default()) + } + ErrorPriority::Ordinary => region_error.not_leader = Some(Default::default()), + } + kvrpcpb::GetResponse { + region_error: Some(region_error), + error: Some(kvrpcpb::KeyError { + locked: Some(kvrpcpb::LockInfo { + key: vec![1], + primary_lock: vec![1], + lock_version: 1, + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + } + } + + #[rstest::rstest] + #[case::ordinary(Some(ErrorPriority::Ordinary), false, true)] + #[case::incompatible(Some(ErrorPriority::Incompatible), false, true)] + #[case::undetermined(Some(ErrorPriority::Undetermined), false, true)] + #[case::resolution_disabled(Some(ErrorPriority::Undetermined), false, false)] + #[case::replay_region_error(Some(ErrorPriority::Undetermined), true, true)] + #[case::replay_success(None, true, true)] + #[tokio::test] + async fn resolve_lock_checks_initial_and_replayed_responses( + #[case] priority: Option, + #[case] replay: bool, + #[case] resolve_enabled: bool, + ) { + let final_response = + priority + .map(mixed_lock_response) + .unwrap_or_else(|| kvrpcpb::GetResponse { + value: b"value".to_vec(), + ..Default::default() + }); + let mut responses = Vec::new(); + if replay { + let mut locked = mixed_lock_response(ErrorPriority::Undetermined); + locked.region_error = None; + responses.push(locked); + } + responses.push(final_response.clone()); + let attempts = Arc::new(AtomicUsize::new(0)); + let auxiliary = Arc::new(std::sync::Mutex::new(Vec::new())); + let recorded = auxiliary.clone(); + let client = crate::mock::MockKvClient::with_dispatch_hook(move |request| { + assert!(replay, "region errors must bypass auxiliary lock RPCs"); + if request.is::() { + recorded.lock().unwrap().push("check"); + Ok(Box::new(kvrpcpb::CheckTxnStatusResponse::default()) as Box) + } else { + assert!(request.is::()); + recorded.lock().unwrap().push("resolve"); + Ok(Box::new(kvrpcpb::ResolveLockResponse::default()) as Box) + } + }); + let plan = ResolveLock { + inner: LockResponsePlan { + responses, + attempts: attempts.clone(), + }, + timestamp: Timestamp::default(), + pd_client: Arc::new(MockPdClient::new(client)), + backoff: if resolve_enabled { + Backoff::no_jitter_backoff(0, 0, 1) + } else { + Backoff::no_backoff() + }, + context: PlanContext::new(Keyspace::Disable), + }; + assert_eq!(plan.execute().await.unwrap(), final_response); + assert_eq!(attempts.load(Ordering::SeqCst), if replay { 2 } else { 1 }); + assert_eq!( + *auxiliary.lock().unwrap(), + if replay { + vec!["check", "resolve"] + } else { + vec![] + } + ); + } + + #[rstest::rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn cleanup_locks_preserves_resolve_lock_errors(#[case] region_error: bool) { + let client = crate::mock::MockKvClient::with_dispatch_hook(move |request| { + if request.is::() { + return Ok(Box::new(kvrpcpb::ScanLockResponse { + locks: vec![kvrpcpb::LockInfo { + key: vec![1], + primary_lock: vec![1], + lock_version: 1, + ..Default::default() + }], + ..Default::default() + })); + } + if request.is::() { + return Ok(Box::new(kvrpcpb::CheckTxnStatusResponse::default())); + } + assert!(request.is::()); + Ok(Box::new(kvrpcpb::ResolveLockResponse { + region_error: region_error.then_some(errorpb::Error { + message: "resolve lock region error".into(), + ..Default::default() + }), + error: (!region_error).then_some(kvrpcpb::KeyError { + abort: "resolve lock key error".into(), + ..Default::default() + }), + ..Default::default() + })) + }); + let plan = crate::request::PlanBuilder::new( + Arc::new(MockPdClient::new(client)), + Keyspace::Disable, + kvrpcpb::ScanLockRequest { + start_key: vec![1], + end_key: vec![2], + max_version: 10, + limit: 1024, + ..Default::default() + }, + ) + .cleanup_locks( + ResolveLocksContext::default(), + ResolveLocksOptions::default(), + Backoff::no_backoff(), + ) + .retry_multi_region(Backoff::no_backoff()) + .plan(); + + match plan.execute().await { + Err(Error::RegionError(error)) if region_error => { + assert_eq!(error.message, "resolve lock region error"); + } + Ok(mut results) if !region_error => { + assert_eq!(results.len(), 1); + let Err(Error::MultipleKeyErrors(errors)) = results.pop().unwrap() else { + panic!("expected the original key error"); + }; + assert_eq!(errors.len(), 1); + assert!(matches!(&errors[0], Error::KeyError(error) + if error.abort == "resolve lock key error")); + } + _ => panic!("expected the original resolve lock error"), + } + } + + #[async_trait] + impl Plan for RegionErrorsPlan { + type Result = Vec; + + async fn execute(&self) -> Result { + Ok(self.0.clone()) + } + } + + #[rstest::rstest] + #[case(ErrorPriority::Undetermined)] + #[case(ErrorPriority::Incompatible)] + #[case(ErrorPriority::Ordinary)] + #[tokio::test] + async fn extract_error_preserves_priority_and_ordinary_wrapping( + #[case] priority: ErrorPriority, + ) { + let mut errors = vec![errorpb::Error { + message: "first ordinary error".into(), + server_is_busy: Some(Default::default()), + ..Default::default() + }]; + if priority != ErrorPriority::Ordinary { + errors.push(errorpb::Error { + incompatible_request: Some(crate::mock::upper_admission_rejection()), + ..Default::default() + }); + } + if priority == ErrorPriority::Undetermined { + errors.push(errorpb::Error { + undetermined_result: Some(Default::default()), + ..Default::default() + }); + } + let result = ExtractError { + inner: RegionErrorsPlan( + errors + .into_iter() + .map(|error| kvrpcpb::GetResponse { + region_error: Some(error), + ..Default::default() + }) + .collect(), + ), + } + .execute() + .await + .unwrap_err(); + match priority { + ErrorPriority::Ordinary => assert!(matches!(result, Error::ExtractedErrors(errors) + if errors.len() == 1 && matches!(&errors[0], Error::RegionError(error) if error.message == "first ordinary error"))), + _ => assert_eq!(result.priority(), priority), + } + } + + #[tokio::test] + async fn all_store_raw_request_does_not_apply_transaction_metadata_or_reload() { + let attempts = Arc::new(AtomicUsize::new(0)); + let recorded = attempts.clone(); + let incompatible = crate::mock::upper_admission_rejection(); + let rejection = incompatible.clone(); + let client = crate::mock::MockKvClient::with_dispatch_hook(move |request| { + let request = request + .downcast_ref::() + .unwrap(); + let context = request.context.as_ref().unwrap(); + assert_eq!(context.txn_protocol_version, 0); + assert_eq!( + context.request_origin, + kvrpcpb::RequestOrigin::Unknown as i32 + ); + recorded.fetch_add(1, Ordering::SeqCst); + Ok(Box::new(kvrpcpb::UnsafeDestroyRangeResponse { + region_error: Some(errorpb::Error { + incompatible_request: Some(rejection.clone()), + ..Default::default() + }), + ..Default::default() + }) as Box) + }); + let fixture = + crate::mock::ProtocolTestClient::new(client, crate::mock::ReloadOutcome::Changed); + let result = crate::request::PlanBuilder::new_with_context( + fixture.pd_client, + kvrpcpb::UnsafeDestroyRangeRequest::default(), + PlanContext::new(Keyspace::Disable) + .with_request_origin(crate::config::RequestOrigin::TiFlash), + ) + .all_stores(Backoff::no_jitter_backoff(0, 0, 1)) + .merge(crate::request::Collect) + .plan() + .execute() + .await; + assert!(matches!(result, Err(Error::IncompatibleRequest(error)) if *error == incompatible)); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + assert!(fixture.reloads.lock().unwrap().is_empty()); + } + #[test] fn grpc_statuses_are_classified_without_cache_churn_for_deadlines() { #[derive(Debug)] @@ -1692,9 +2151,10 @@ mod test { timestamp: Timestamp::default(), backoff: Backoff::no_backoff(), pd_client: Arc::new(MockPdClient::default()), - keyspace: Keyspace::Disable, + context: PlanContext::new(Keyspace::Disable), }, pd_client: Arc::new(MockPdClient::default()), + context: PlanContext::new(Keyspace::Disable), backoff: Backoff::no_backoff(), preserve_region_results: false, terminal_on_undetermined: false, @@ -1774,6 +2234,7 @@ mod test { let plan = RetryableMultiRegion { inner, pd_client: Arc::new(MockPdClient::default()), + context: PlanContext::new(Keyspace::Disable), backoff: Backoff::no_jitter_backoff(1, 2, 10), preserve_region_results: false, terminal_on_undetermined: true, @@ -1804,6 +2265,7 @@ mod test { let plan = RetryableMultiRegion { inner, pd_client: Arc::new(MockPdClient::default()), + context: PlanContext::new(Keyspace::Disable), backoff: Backoff::no_jitter_backoff(1, 2, RETRIES), preserve_region_results: false, terminal_on_undetermined: false, @@ -1867,6 +2329,7 @@ mod test { let plan = RetryableMultiRegion { inner, pd_client: Arc::new(MockPdClient::default()), + context: PlanContext::new(Keyspace::Disable), backoff: Backoff::no_jitter_backoff(1, 2, 10), preserve_region_results: false, terminal_on_undetermined: true, @@ -1897,6 +2360,7 @@ mod test { let plan = RetryableMultiRegion { inner, pd_client: Arc::new(MockPdClient::default()), + context: PlanContext::new(Keyspace::Disable), backoff: Backoff::no_jitter_backoff(1, 2, RETRIES), preserve_region_results: false, terminal_on_undetermined: true, @@ -1925,6 +2389,13 @@ mod test { })); assert!(is_undetermined_region_error(&undetermined)); + let both = Error::from(errorpb::Error { + undetermined_result: Some(errorpb::UndeterminedResult::default()), + incompatible_request: Some(errorpb::IncompatibleRequest::default()), + ..Default::default() + }); + assert!(is_undetermined_region_error(&both)); + let busy = Error::RegionError(Box::new(errorpb::Error { server_is_busy: Some(errorpb::ServerIsBusy::default()), ..Default::default() @@ -1997,6 +2468,7 @@ mod test { let plan = RetryableMultiRegion { inner: MaskingPlan { idx: 0 }, pd_client: Arc::new(MockPdClient::default()), + context: PlanContext::new(Keyspace::Disable), backoff: Backoff::no_backoff(), preserve_region_results: false, terminal_on_undetermined: true, diff --git a/src/request/plan_builder.rs b/src/request/plan_builder.rs index 183c3e6b..b8809e96 100644 --- a/src/request/plan_builder.rs +++ b/src/request/plan_builder.rs @@ -6,6 +6,8 @@ use std::sync::Arc; use super::plan::PreserveShard; use super::Keyspace; use crate::backoff::Backoff; +use crate::config::RequestOrigin; +use crate::config::DEFAULT_TXN_PROTOCOL_VERSION; use crate::pd::PdClient; use crate::request::plan::{CleanupLocks, RetryableAllStores}; use crate::request::shard::HasNextBatch; @@ -35,10 +37,63 @@ use crate::Timestamp; /// Builder type for plans (see that module for more). pub struct PlanBuilder { pd_client: Arc, + context: PlanContext, plan: P, phantom: PhantomData, } +/// Immutable client state required to prepare and execute a request plan. +/// +/// Transaction protocol metadata is ignored for requests classified as raw. +#[derive(Clone, Copy, Debug)] +pub(crate) struct PlanContext { + keyspace: Keyspace, + request_origin: RequestOrigin, + default_txn_protocol_version: u32, +} + +impl PlanContext { + pub(crate) fn new(keyspace: Keyspace) -> Self { + Self { + keyspace, + request_origin: RequestOrigin::Unknown, + default_txn_protocol_version: DEFAULT_TXN_PROTOCOL_VERSION, + } + } + + #[must_use] + pub(crate) fn with_keyspace(mut self, keyspace: Keyspace) -> Self { + self.keyspace = keyspace; + self + } + + /// Set the process origin attached to transaction RPCs. + #[must_use] + pub(crate) fn with_request_origin(mut self, request_origin: RequestOrigin) -> Self { + self.request_origin = request_origin; + self + } + + /// Set the highest transaction protocol version this client can understand. + #[must_use] + pub(crate) fn with_default_txn_protocol_version(mut self, version: u32) -> Self { + self.default_txn_protocol_version = version; + self + } + + pub(crate) fn keyspace(&self) -> Keyspace { + self.keyspace + } + + pub(crate) fn default_txn_protocol_version(&self) -> u32 { + self.default_txn_protocol_version + } + + pub(crate) fn request_origin(&self) -> RequestOrigin { + self.request_origin + } +} + /// Used to ensure that a plan has a designated target or targets, a target is /// a particular TiKV server. pub trait PlanBuilderPhase {} @@ -48,13 +103,23 @@ pub struct Targetted; impl PlanBuilderPhase for Targetted {} impl PlanBuilder, NoTarget> { - pub fn new(pd_client: Arc, keyspace: Keyspace, mut request: Req) -> Self { - request.set_api_version(keyspace.api_version()); + pub fn new(pd_client: Arc, keyspace: Keyspace, request: Req) -> Self { + Self::new_with_context(pd_client, request, PlanContext::new(keyspace)) + } + + pub(crate) fn new_with_context( + pd_client: Arc, + mut request: Req, + context: PlanContext, + ) -> Self { + request.set_api_version(context.keyspace.api_version()); PlanBuilder { - pd_client, + pd_client: pd_client.clone(), + context, plan: Dispatch { request, kv_client: None, + context, }, phantom: PhantomData, } @@ -75,19 +140,19 @@ impl PlanBuilder { self, timestamp: Timestamp, backoff: Backoff, - keyspace: Keyspace, ) -> PlanBuilder, Ph> where - P::Result: HasLocks, + P::Result: HasLocks + HasRegionError, { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: ResolveLock { inner: self.plan, timestamp, backoff, pd_client: self.pd_client, - keyspace, + context: self.context, }, phantom: PhantomData, } @@ -98,7 +163,6 @@ impl PlanBuilder { ctx: ResolveLocksContext, options: ResolveLocksOptions, backoff: Backoff, - keyspace: Keyspace, ) -> PlanBuilder, Ph> where P: Shardable + NextBatch, @@ -106,6 +170,7 @@ impl PlanBuilder { { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: CleanupLocks { inner: self.plan, ctx, @@ -113,7 +178,7 @@ impl PlanBuilder { store: None, backoff, pd_client: self.pd_client, - keyspace, + context: self.context, }, phantom: PhantomData, } @@ -128,6 +193,7 @@ impl PlanBuilder { { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: MergeResponse { inner: self.plan, merge, @@ -147,6 +213,7 @@ impl PlanBuilder { { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: ProcessResponse { inner: self.plan, processor: DefaultProcessor, @@ -235,9 +302,11 @@ where ) -> PlanBuilder, Targetted> { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: RetryableMultiRegion { inner: self.plan, pd_client: self.pd_client, + context: self.context, backoff, preserve_region_results: flags.preserve_region_results, terminal_on_undetermined: flags.terminal_on_undetermined, @@ -254,7 +323,7 @@ impl PlanBuilder, NoTarget> { self, store: RegionStore, ) -> Result, Targetted>> { - set_single_region_store(self.plan, store, self.pd_client) + set_single_region_store(self.plan, store, self.pd_client, self.context) } } @@ -268,9 +337,11 @@ where ) -> PlanBuilder, Targetted> { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: RetryableAllStores { inner: self.plan, pd_client: self.pd_client, + context: self.context, backoff, }, phantom: PhantomData, @@ -285,6 +356,7 @@ where pub fn preserve_shard(self) -> PlanBuilder, NoTarget> { PlanBuilder { pd_client: self.pd_client.clone(), + context: self.context, plan: PreserveShard { inner: self.plan, shard: None, @@ -301,6 +373,7 @@ where pub fn extract_error(self) -> PlanBuilder, Targetted> { PlanBuilder { pd_client: self.pd_client, + context: self.context, plan: ExtractError { inner: self.plan }, phantom: self.phantom, } @@ -311,12 +384,19 @@ fn set_single_region_store( mut plan: Dispatch, store: RegionStore, pd_client: Arc, + context: PlanContext, ) -> Result, Targetted>> { + plan.request.prepare_txn_rpc( + store.txn_protocol_version_range, + context.default_txn_protocol_version(), + context.request_origin().as_proto(), + )?; plan.request.set_leader(&store.region_with_leader)?; plan.kv_client = Some(store.client); Ok(PlanBuilder { plan, pd_client, + context, phantom: PhantomData, }) } @@ -326,3 +406,46 @@ pub trait SingleKey { #[allow(clippy::ptr_arg)] fn key(&self) -> &Vec; } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::config::RequestOrigin; + use crate::mock::{MockKvClient, MockPdClient}; + use crate::proto::kvrpcpb; + use crate::store::TxnProtocolVersionRange; + + #[tokio::test] + async fn context_metadata_is_applied_to_single_region_transaction_requests() { + let context = PlanContext::new(Keyspace::Disable) + .with_request_origin(RequestOrigin::TiFlash) + .with_default_txn_protocol_version(1); + let request = kvrpcpb::GetRequest { + context: Some(kvrpcpb::Context::default()), + ..Default::default() + }; + let store = RegionStore::with_metadata( + MockPdClient::region1(), + 1, + TxnProtocolVersionRange { min: 0, max: 2 }, + Arc::new(MockKvClient::default()), + ); + + let builder = PlanBuilder::new_with_context( + Arc::new(MockPdClient::new(MockKvClient::default())), + request, + context, + ) + .single_region_with_store(store) + .await + .unwrap(); + let request_context = builder.plan.request.context.as_ref().unwrap(); + assert_eq!(request_context.txn_protocol_version, 1); + assert_eq!( + request_context.request_origin, + kvrpcpb::RequestOrigin::TiFlash as i32 + ); + } +} diff --git a/src/request/shard.rs b/src/request/shard.rs index 1bac69e4..6901f075 100644 --- a/src/request/shard.rs +++ b/src/request/shard.rs @@ -14,6 +14,7 @@ use crate::request::Plan; use crate::request::ResolveLock; use crate::store::RegionStore; use crate::store::Request; +use crate::store::TxnProtocolRequirement; use crate::Result; use std::fmt::Debug; @@ -35,6 +36,10 @@ macro_rules! impl_inner_shardable { fn apply_store(&mut self, store: &RegionStore) -> Result<()> { self.inner.apply_store(store) } + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + self.inner.txn_protocol_requirement() + } }; } @@ -59,6 +64,10 @@ pub trait Shardable { } fn apply_store(&mut self, store: &RegionStore) -> Result<()>; + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + TxnProtocolRequirement::NotTransaction + } } pub trait Batchable { @@ -120,13 +129,23 @@ impl Shardable for Dispatch { Dispatch { request: self.request.clone_then_apply_shard(shard), kv_client: self.kv_client.clone(), + context: self.context, } } fn apply_store(&mut self, store: &RegionStore) -> Result<()> { self.kv_client = Some(store.client.clone()); + self.request.prepare_txn_rpc( + store.txn_protocol_version_range, + self.context.default_txn_protocol_version(), + self.context.request_origin().as_proto(), + )?; self.request.apply_store(store) } + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + Request::txn_protocol_requirement(&self.request) + } } impl NextBatch for Dispatch { @@ -153,6 +172,10 @@ impl Shardable for PreserveShard

{ fn apply_store(&mut self, store: &RegionStore) -> Result<()> { self.inner.apply_store(store) } + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + self.inner.txn_protocol_requirement() + } } impl Shardable for ResolveLock { @@ -177,6 +200,10 @@ impl Shardable for CleanupLocks { self.store = Some(store.clone()); self.inner.apply_store(store) } + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + self.inner.txn_protocol_requirement() + } } #[doc(hidden)] @@ -207,6 +234,10 @@ macro_rules! shardable_key { fn apply_store(&mut self, store: &$crate::store::RegionStore) -> $crate::Result<()> { self.set_leader(&store.region_with_leader) } + + fn txn_protocol_requirement(&self) -> $crate::store::TxnProtocolRequirement { + $crate::store::Request::txn_protocol_requirement(self) + } } }; } @@ -237,6 +268,10 @@ macro_rules! shardable_keys { fn apply_store(&mut self, store: &$crate::store::RegionStore) -> $crate::Result<()> { self.set_leader(&store.region_with_leader) } + + fn txn_protocol_requirement(&self) -> $crate::store::TxnProtocolRequirement { + $crate::store::Request::txn_protocol_requirement(self) + } } }; } @@ -302,6 +337,10 @@ macro_rules! shardable_range { fn apply_store(&mut self, store: &$crate::store::RegionStore) -> $crate::Result<()> { self.set_leader(&store.region_with_leader) } + + fn txn_protocol_requirement(&self) -> $crate::store::TxnProtocolRequirement { + $crate::store::Request::txn_protocol_requirement(self) + } } }; } diff --git a/src/store/errors.rs b/src/store/errors.rs index 06a49063..b91dd168 100644 --- a/src/store/errors.rs +++ b/src/store/errors.rs @@ -5,6 +5,9 @@ use crate::Error; // Those that can have a single region error pub trait HasRegionError { + /// Check for a region error without consuming or modifying the response. + fn has_region_error(&self) -> bool; + fn region_error(&mut self) -> Option; } @@ -26,6 +29,10 @@ impl HasRegionErrors for T { macro_rules! has_region_error { ($type:ty) => { impl HasRegionError for $type { + fn has_region_error(&self) -> bool { + self.region_error.is_some() + } + fn region_error(&mut self) -> Option { self.region_error.take().map(|e| e.into()) } @@ -186,6 +193,11 @@ impl HasKeyErrors for Vec { } impl HasRegionError for Result { + fn has_region_error(&self) -> bool { + self.as_ref() + .is_ok_and(|response| response.has_region_error()) + } + fn region_error(&mut self) -> Option { self.as_mut().ok().and_then(|t| t.region_error()) } @@ -216,9 +228,34 @@ fn extract_errors( #[cfg(test)] mod test { use super::HasKeyErrors; + use super::HasRegionError; use crate::common::Error; use crate::internal_err; use crate::proto::kvrpcpb; + + #[test] + fn region_error_query_is_non_destructive_and_forwarded() { + let error = crate::proto::errorpb::Error { + message: "preserve me".into(), + ..Default::default() + }; + let response = kvrpcpb::GetResponse { + region_error: Some(error.clone()), + value: b"value".to_vec(), + ..Default::default() + }; + let mut wrapped = + crate::request::plan::ResponseWithShard(Ok::<_, Error>(response.clone()), 7); + assert!(wrapped.has_region_error()); + assert!(wrapped.has_region_error()); + assert_eq!(wrapped.0.as_ref().unwrap(), &response); + assert_eq!(wrapped.region_error(), Some(error)); + assert!(!wrapped.has_region_error()); + assert_eq!(wrapped.0.as_ref().unwrap().value, b"value"); + + let failed: Result = Err(Error::Unimplemented); + assert!(!failed.has_region_error()); + } #[test] fn result_haslocks() { let mut resp: Result<_, Error> = Ok(kvrpcpb::CommitResponse::default()); diff --git a/src/store/mod.rs b/src/store/mod.rs index fb6ed709..cb232162 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -8,7 +8,6 @@ use std::cmp::max; use std::cmp::min; use std::sync::Arc; -use derive_new::new; use futures::prelude::*; use futures::stream::BoxStream; @@ -18,6 +17,7 @@ pub use self::client::TikvConnect; pub use self::errors::HasKeyErrors; pub use self::errors::HasRegionError; pub use self::errors::HasRegionErrors; +pub(crate) use self::request::reload_txn_protocol_range; pub use self::request::Request; use crate::pd::PdClient; use crate::proto::kvrpcpb; @@ -26,17 +26,122 @@ use crate::BoundRange; use crate::Key; use crate::Result; -#[derive(new, Clone)] +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct TxnProtocolVersionRange { + pub min: u32, + pub max: u32, +} + +impl TxnProtocolVersionRange { + /// The version the client would send before validating the store's range. + pub(crate) fn preferred_version(self, default_version: u32) -> u32 { + default_version.min(self.max) + } + + /// Select a version supported by both the request and the store. + pub(crate) fn select(self, default_version: u32, required_version: u32) -> Option { + let selected = self.preferred_version(default_version); + (self.min <= self.max && selected >= self.min && selected >= required_version) + .then_some(selected) + } +} + +impl From<&crate::proto::metapb::Store> for TxnProtocolVersionRange { + fn from(store: &crate::proto::metapb::Store) -> Self { + store + .txn_protocol_version_range + .as_ref() + .map(|range| Self { + min: range.min, + max: range.max, + }) + .unwrap_or_default() + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum TxnProtocolRequirement { + NotTransaction, + Transaction { required_version: u32 }, +} + +impl TxnProtocolRequirement { + pub(crate) fn select( + self, + range: TxnProtocolVersionRange, + default_version: u32, + ) -> Option { + match self { + Self::NotTransaction => None, + Self::Transaction { required_version } => { + range.select(default_version, required_version) + } + } + } +} + +#[derive(Clone)] pub struct RegionStore { pub region_with_leader: RegionWithLeader, + pub store_id: u64, + pub txn_protocol_version_range: TxnProtocolVersionRange, pub client: Arc, } -#[derive(new, Clone)] +impl RegionStore { + pub fn new( + region_with_leader: RegionWithLeader, + client: Arc, + ) -> Self { + let store_id = region_with_leader.get_store_id().unwrap_or_default(); + Self::with_metadata( + region_with_leader, + store_id, + TxnProtocolVersionRange::default(), + client, + ) + } + + pub fn with_metadata( + region_with_leader: RegionWithLeader, + store_id: u64, + txn_protocol_version_range: TxnProtocolVersionRange, + client: Arc, + ) -> Self { + Self { + region_with_leader, + store_id, + txn_protocol_version_range, + client, + } + } +} + +#[derive(Clone)] pub struct Store { + pub store_id: u64, + pub txn_protocol_version_range: TxnProtocolVersionRange, pub client: Arc, } +impl Store { + pub fn new(client: Arc) -> Self { + Self::with_metadata(0, TxnProtocolVersionRange::default(), client) + } + + pub fn with_metadata( + store_id: u64, + txn_protocol_version_range: TxnProtocolVersionRange, + client: Arc, + ) -> Self { + Self { + store_id, + txn_protocol_version_range, + client, + } + } +} + /// Maps keys to a stream of stores. `key_data` must be sorted in increasing order pub fn region_stream_for_keys( key_data: impl Iterator + Send + Sync + 'static, @@ -92,3 +197,38 @@ pub fn region_stream_for_ranges( ) -> BoxStream<'static, Result<(Vec, RegionWithLeader)>> { pd_client.clone().group_ranges_by_region(ranges) } + +#[cfg(test)] +mod txn_protocol_tests { + use super::*; + + #[rstest::rstest] + #[case(None, TxnProtocolVersionRange { min: 0, max: 0 })] + #[case(Some(crate::proto::metapb::TxnProtocolVersionRange { min: 2, max: 1 }), TxnProtocolVersionRange { min: 2, max: 1 })] + fn store_range_conversion_preserves_invalid_metadata( + #[case] range: Option, + #[case] expected: TxnProtocolVersionRange, + ) { + let store = crate::proto::metapb::Store { + txn_protocol_version_range: range, + ..Default::default() + }; + assert_eq!(TxnProtocolVersionRange::from(&store), expected); + } + + #[test] + fn select_txn_protocol_version_validates_the_whole_range() { + let range = TxnProtocolVersionRange { min: 0, max: 2 }; + assert_eq!(range.select(1, 0), Some(1)); + assert_eq!(range.select(1, 2), None); + assert_eq!(range.select(2, 2), Some(2)); + assert_eq!( + TxnProtocolVersionRange { min: 2, max: 1 }.select(2, 0), + None + ); + assert_eq!( + TxnProtocolVersionRange { min: 2, max: 3 }.select(1, 0), + None + ); + } +} diff --git a/src/store/request.rs b/src/store/request.rs index 2d9c29ad..a76984a3 100644 --- a/src/store/request.rs +++ b/src/store/request.rs @@ -1,15 +1,22 @@ // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0. use std::any::Any; +use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; use tonic::transport::Channel; use tonic::IntoRequest; +use crate::backoff::Backoff; +use crate::pd::PdClient; +use crate::proto::errorpb; use crate::proto::kvrpcpb; use crate::proto::tikvpb::tikv_client::TikvClient; +use crate::region::StoreId; use crate::store::RegionWithLeader; +use crate::store::TxnProtocolRequirement; +use crate::store::TxnProtocolVersionRange; use crate::Error; use crate::Result; @@ -24,10 +31,109 @@ pub trait Request: Any + Sync + Send + 'static { fn as_any(&self) -> &dyn Any; fn set_leader(&mut self, leader: &RegionWithLeader) -> Result<()>; fn set_api_version(&mut self, api_version: kvrpcpb::ApiVersion); + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement; + /// Validate and select the transaction protocol version, then apply RPC metadata. + /// Incompatible requests leave the context unchanged; non-transactional requests + /// are untouched. Routing and leader selection are handled separately. + fn prepare_txn_rpc( + &mut self, + range: TxnProtocolVersionRange, + default_version: u32, + request_origin: i32, + ) -> Result<()>; +} + +fn select_txn_protocol_version( + range: TxnProtocolVersionRange, + default_version: u32, + required_version: u32, +) -> Result { + if let Some(selected) = range.select(default_version, required_version) { + return Ok(selected); + } + let selected = range.preferred_version(default_version); + let incompatible = crate::proto::errorpb::IncompatibleRequest { + // This is a client-side error, not a server admission rejection. + // `Unknown` keeps the reload/retry path from misclassifying it. + reason: crate::proto::errorpb::IncompatibleRequestReason::Unknown as i32, + message: format!( + "no compatible transaction protocol version: client={default_version}, required={required_version}, store=[{}, {}]", + range.min, range.max + ), + provided_txn_protocol_version: selected, + min_compatible_txn_protocol_version: range.min, + max_compatible_txn_protocol_version: range.max, + }; + Err(Error::IncompatibleRequest(Box::new(incompatible))) +} + +fn is_upper_txn_protocol_admission_rejection( + incompatible: &errorpb::IncompatibleRequest, + selected_version: u32, +) -> bool { + incompatible.reason == errorpb::IncompatibleRequestReason::TxnProtocolVersionOutOfRange as i32 + && incompatible.min_compatible_txn_protocol_version + <= incompatible.max_compatible_txn_protocol_version + && incompatible.provided_txn_protocol_version == selected_version + && incompatible.provided_txn_protocol_version + > incompatible.max_compatible_txn_protocol_version +} + +/// Reload only after a side-effect-free upper admission rejection, and only +/// propose a retry if the refreshed metadata selects a different valid version. +/// Callers retain responsibility for routing, applying the range, and returning +/// the original incompatible error when no progress is possible. +pub(crate) async fn reload_txn_protocol_range( + pd_client: Arc, + store_id: StoreId, + old_range: TxnProtocolVersionRange, + default_version: u32, + requirement: TxnProtocolRequirement, + incompatible: &errorpb::IncompatibleRequest, + backoff: &mut Backoff, +) -> Option { + let selected = requirement.select(old_range, default_version)?; + if !is_upper_txn_protocol_admission_rejection(incompatible, selected) + || backoff.next_delay_duration().is_none() + { + return None; + } + let range = pd_client.reload_store(store_id).await.ok()??; + let new_version = requirement.select(range, default_version)?; + (new_version != selected).then_some(range) +} + +// Derive the version from the payload so neither shared-lock op can be +// declared as legacy, even when it appears in the wrong request type. +fn required_txn_protocol_version(mutations: &[kvrpcpb::Mutation]) -> u32 { + if mutations.iter().any(|mutation| { + mutation.op == kvrpcpb::Op::SharedLock as i32 + || mutation.op == kvrpcpb::Op::SharedPessimisticLock as i32 + }) { + kvrpcpb::TxnProtocolVersion::TxnVerSupportSharedLock as u32 + } else { + 0 + } } macro_rules! impl_request { + ($name: ident, $fun: ident, $label: literal, $requirement: expr) => { + impl_request_with_requirement!($name, $fun, $label, |_: &kvrpcpb::$name| $requirement); + }; +} + +macro_rules! impl_txn_request_with_payload_requirement { ($name: ident, $fun: ident, $label: literal) => { + impl_request_with_requirement!($name, $fun, $label, |request: &kvrpcpb::$name| { + TxnProtocolRequirement::Transaction { + required_version: required_txn_protocol_version(&request.mutations), + } + }); + }; +} + +macro_rules! impl_request_with_requirement { + ($name: ident, $fun: ident, $label: literal, $classifier: expr) => { #[async_trait] impl Request for kvrpcpb::$name { async fn dispatch( @@ -68,56 +174,382 @@ macro_rules! impl_request { let ctx = self.context.get_or_insert(kvrpcpb::Context::default()); ctx.api_version = api_version.into(); } + + fn txn_protocol_requirement(&self) -> TxnProtocolRequirement { + ($classifier)(self) + } + + fn prepare_txn_rpc( + &mut self, + range: TxnProtocolVersionRange, + default_version: u32, + request_origin: i32, + ) -> Result<()> { + match self.txn_protocol_requirement() { + TxnProtocolRequirement::NotTransaction => Ok(()), + TxnProtocolRequirement::Transaction { required_version } => { + let selected = + select_txn_protocol_version(range, default_version, required_version)?; + let context = self.context.get_or_insert_with(kvrpcpb::Context::default); + context.txn_protocol_version = selected; + if context.request_origin == kvrpcpb::RequestOrigin::Unknown as i32 { + context.request_origin = request_origin; + } + Ok(()) + } + } + } } }; } -impl_request!(RawGetRequest, raw_get, "raw_get"); -impl_request!(RawBatchGetRequest, raw_batch_get, "raw_batch_get"); -impl_request!(RawGetKeyTtlRequest, raw_get_key_ttl, "raw_get_key_ttl"); -impl_request!(RawPutRequest, raw_put, "raw_put"); -impl_request!(RawBatchPutRequest, raw_batch_put, "raw_batch_put"); -impl_request!(RawDeleteRequest, raw_delete, "raw_delete"); -impl_request!(RawBatchDeleteRequest, raw_batch_delete, "raw_batch_delete"); -impl_request!(RawScanRequest, raw_scan, "raw_scan"); -impl_request!(RawBatchScanRequest, raw_batch_scan, "raw_batch_scan"); -impl_request!(RawDeleteRangeRequest, raw_delete_range, "raw_delete_range"); -impl_request!(RawCasRequest, raw_compare_and_swap, "raw_compare_and_swap"); -impl_request!(RawCoprocessorRequest, raw_coprocessor, "raw_coprocessor"); - -impl_request!(GetRequest, kv_get, "kv_get"); -impl_request!(ScanRequest, kv_scan, "kv_scan"); -impl_request!(PrewriteRequest, kv_prewrite, "kv_prewrite"); -impl_request!(CommitRequest, kv_commit, "kv_commit"); -impl_request!(BatchGetRequest, kv_batch_get, "kv_batch_get"); -impl_request!(BatchRollbackRequest, kv_batch_rollback, "kv_batch_rollback"); -impl_request!( +macro_rules! impl_raw_request { + ($name: ident, $fun: ident, $label: literal) => { + impl_request!($name, $fun, $label, TxnProtocolRequirement::NotTransaction); + }; +} + +macro_rules! impl_txn_request { + ($name: ident, $fun: ident, $label: literal) => { + impl_request!( + $name, + $fun, + $label, + (TxnProtocolRequirement::Transaction { + required_version: 0 + }) + ); + }; +} + +impl_raw_request!(RawGetRequest, raw_get, "raw_get"); +impl_raw_request!(RawBatchGetRequest, raw_batch_get, "raw_batch_get"); +impl_raw_request!(RawGetKeyTtlRequest, raw_get_key_ttl, "raw_get_key_ttl"); +impl_raw_request!(RawPutRequest, raw_put, "raw_put"); +impl_raw_request!(RawBatchPutRequest, raw_batch_put, "raw_batch_put"); +impl_raw_request!(RawDeleteRequest, raw_delete, "raw_delete"); +impl_raw_request!(RawBatchDeleteRequest, raw_batch_delete, "raw_batch_delete"); +impl_raw_request!(RawScanRequest, raw_scan, "raw_scan"); +impl_raw_request!(RawBatchScanRequest, raw_batch_scan, "raw_batch_scan"); +impl_raw_request!(RawDeleteRangeRequest, raw_delete_range, "raw_delete_range"); +impl_raw_request!(RawCasRequest, raw_compare_and_swap, "raw_compare_and_swap"); +impl_raw_request!(RawCoprocessorRequest, raw_coprocessor, "raw_coprocessor"); + +impl_txn_request!(GetRequest, kv_get, "kv_get"); +impl_txn_request!(ScanRequest, kv_scan, "kv_scan"); +impl_txn_request_with_payload_requirement!(PrewriteRequest, kv_prewrite, "kv_prewrite"); +impl_txn_request!(CommitRequest, kv_commit, "kv_commit"); +impl_txn_request!(BatchGetRequest, kv_batch_get, "kv_batch_get"); +impl_txn_request!(BatchRollbackRequest, kv_batch_rollback, "kv_batch_rollback"); +impl_txn_request!( PessimisticRollbackRequest, kv_pessimistic_rollback, "kv_pessimistic_rollback" ); -impl_request!(ResolveLockRequest, kv_resolve_lock, "kv_resolve_lock"); -impl_request!(ScanLockRequest, kv_scan_lock, "kv_scan_lock"); -impl_request!( +impl_txn_request!(ResolveLockRequest, kv_resolve_lock, "kv_resolve_lock"); +impl_txn_request!(ScanLockRequest, kv_scan_lock, "kv_scan_lock"); +impl_txn_request_with_payload_requirement!( PessimisticLockRequest, kv_pessimistic_lock, "kv_pessimistic_lock" ); -impl_request!(TxnHeartBeatRequest, kv_txn_heart_beat, "kv_txn_heart_beat"); -impl_request!( +impl_txn_request!(TxnHeartBeatRequest, kv_txn_heart_beat, "kv_txn_heart_beat"); +impl_txn_request!( CheckTxnStatusRequest, kv_check_txn_status, "kv_check_txn_status" ); -impl_request!( +impl_txn_request!( CheckSecondaryLocksRequest, kv_check_secondary_locks, "kv_check_secondary_locks_request" ); -impl_request!(GcRequest, kv_gc, "kv_gc"); -impl_request!(DeleteRangeRequest, kv_delete_range, "kv_delete_range"); +impl_txn_request!(GcRequest, kv_gc, "kv_gc"); +impl_txn_request!(DeleteRangeRequest, kv_delete_range, "kv_delete_range"); impl_request!( UnsafeDestroyRangeRequest, unsafe_destroy_range, - "unsafe_destroy_range" + "unsafe_destroy_range", + TxnProtocolRequirement::NotTransaction ); + +#[cfg(test)] +mod tests { + use super::*; + use crate::mock::{MockKvClient, ProtocolTestClient, ReloadOutcome}; + + #[rstest::rstest] + #[case(ReloadOutcome::Changed, Some(TxnProtocolVersionRange { min: 0, max: 0 }))] + #[case(ReloadOutcome::Unchanged, None)] + #[case(ReloadOutcome::Invalid, None)] + #[case(ReloadOutcome::Missing, None)] + #[case(ReloadOutcome::Failed, None)] + #[tokio::test] + async fn reload_requires_a_different_compatible_version( + #[case] outcome: ReloadOutcome, + #[case] expected: Option, + ) { + let fixture = ProtocolTestClient::new(MockKvClient::default(), outcome); + let range = reload_txn_protocol_range( + fixture.pd_client, + 41, + TxnProtocolVersionRange { min: 0, max: 1 }, + 1, + TxnProtocolRequirement::Transaction { + required_version: 0, + }, + &crate::mock::upper_admission_rejection(), + &mut Backoff::no_jitter_backoff(0, 0, 1), + ) + .await; + assert_eq!(range, expected); + assert_eq!(*fixture.reloads.lock().unwrap(), vec![41]); + } + + #[test] + fn only_a_well_formed_upper_admission_rejection_can_reload() { + let mut incompatible = crate::mock::upper_admission_rejection(); + assert!(is_upper_txn_protocol_admission_rejection(&incompatible, 1)); + incompatible.reason = errorpb::IncompatibleRequestReason::Unknown as i32; + assert!(!is_upper_txn_protocol_admission_rejection(&incompatible, 1)); + incompatible.reason = + errorpb::IncompatibleRequestReason::TxnProtocolVersionOutOfRange as i32; + incompatible.min_compatible_txn_protocol_version = 2; + assert!(!is_upper_txn_protocol_admission_rejection(&incompatible, 1)); + incompatible.min_compatible_txn_protocol_version = 0; + assert!(!is_upper_txn_protocol_admission_rejection(&incompatible, 0)); + } + + #[tokio::test] + async fn reload_does_not_downgrade_below_the_request_requirement() { + let fixture = ProtocolTestClient::new(MockKvClient::default(), ReloadOutcome::Changed); + let incompatible = errorpb::IncompatibleRequest { + provided_txn_protocol_version: 2, + ..crate::mock::upper_admission_rejection() + }; + let range = reload_txn_protocol_range( + fixture.pd_client, + 41, + TxnProtocolVersionRange { min: 0, max: 2 }, + 2, + TxnProtocolRequirement::Transaction { + required_version: 2, + }, + &incompatible, + &mut Backoff::no_jitter_backoff(0, 0, 1), + ) + .await; + assert!(range.is_none()); + assert_eq!(*fixture.reloads.lock().unwrap(), vec![41]); + } + + #[tokio::test] + async fn raw_requests_and_exhausted_budgets_do_not_reload() { + let fixture = ProtocolTestClient::new(MockKvClient::default(), ReloadOutcome::Changed); + for (requirement, mut backoff) in [ + ( + TxnProtocolRequirement::NotTransaction, + Backoff::no_jitter_backoff(0, 0, 1), + ), + ( + TxnProtocolRequirement::Transaction { + required_version: 0, + }, + Backoff::no_backoff(), + ), + ] { + let range = reload_txn_protocol_range( + fixture.pd_client.clone(), + 41, + TxnProtocolVersionRange { min: 0, max: 1 }, + 1, + requirement, + &crate::mock::upper_admission_rejection(), + &mut backoff, + ) + .await; + assert!(range.is_none()); + } + assert!(fixture.reloads.lock().unwrap().is_empty()); + } + + #[test] + fn request_classification_is_explicit() { + macro_rules! assert_txn { + ($($request:ty),+ $(,)?) => {$({ + let request = <$request>::default(); + assert_eq!( + Request::txn_protocol_requirement(&request), + TxnProtocolRequirement::Transaction { required_version: 0 }, + ); + })+}; + } + macro_rules! assert_raw { + ($($request:ty),+ $(,)?) => {$({ + let request = <$request>::default(); + assert_eq!( + Request::txn_protocol_requirement(&request), + TxnProtocolRequirement::NotTransaction, + ); + let mut request = request; + request.prepare_txn_rpc( + TxnProtocolVersionRange { min: 2, max: 1 }, + 1, + kvrpcpb::RequestOrigin::TiFlash as i32, + ).unwrap(); + assert!(request.context.is_none()); + })+}; + } + + assert_raw!( + kvrpcpb::RawGetRequest, + kvrpcpb::RawBatchGetRequest, + kvrpcpb::RawGetKeyTtlRequest, + kvrpcpb::RawPutRequest, + kvrpcpb::RawBatchPutRequest, + kvrpcpb::RawDeleteRequest, + kvrpcpb::RawBatchDeleteRequest, + kvrpcpb::RawScanRequest, + kvrpcpb::RawBatchScanRequest, + kvrpcpb::RawDeleteRangeRequest, + kvrpcpb::RawCasRequest, + kvrpcpb::RawCoprocessorRequest, + kvrpcpb::UnsafeDestroyRangeRequest, + ); + assert_txn!( + kvrpcpb::GetRequest, + kvrpcpb::ScanRequest, + kvrpcpb::PrewriteRequest, + kvrpcpb::CommitRequest, + kvrpcpb::BatchGetRequest, + kvrpcpb::BatchRollbackRequest, + kvrpcpb::PessimisticRollbackRequest, + kvrpcpb::ResolveLockRequest, + kvrpcpb::ScanLockRequest, + kvrpcpb::PessimisticLockRequest, + kvrpcpb::TxnHeartBeatRequest, + kvrpcpb::CheckTxnStatusRequest, + kvrpcpb::CheckSecondaryLocksRequest, + kvrpcpb::GcRequest, + kvrpcpb::DeleteRangeRequest, + ); + } + + #[test] + fn payload_requirement_recognizes_both_shared_ops_in_both_requests() { + fn assert_requirement(mut request: impl Request, required_version: u32) { + assert_eq!( + request.txn_protocol_requirement(), + TxnProtocolRequirement::Transaction { required_version }, + ); + let result = request.prepare_txn_rpc( + TxnProtocolVersionRange { min: 0, max: 1 }, + 1, + kvrpcpb::RequestOrigin::Unknown as i32, + ); + if required_version == 0 { + assert!(result.is_ok()); + } else { + assert!(matches!(result, Err(Error::IncompatibleRequest(_)))); + } + } + + for (op, required_version) in [ + (kvrpcpb::Op::Put, 0), + ( + kvrpcpb::Op::SharedLock, + kvrpcpb::TxnProtocolVersion::TxnVerSupportSharedLock as u32, + ), + ( + kvrpcpb::Op::SharedPessimisticLock, + kvrpcpb::TxnProtocolVersion::TxnVerSupportSharedLock as u32, + ), + ] { + let mutations = vec![ + kvrpcpb::Mutation { + op: kvrpcpb::Op::Put as i32, + ..Default::default() + }, + kvrpcpb::Mutation { + op: op as i32, + ..Default::default() + }, + ]; + assert_requirement( + kvrpcpb::PrewriteRequest { + mutations: mutations.clone(), + ..Default::default() + }, + required_version, + ); + assert_requirement( + kvrpcpb::PessimisticLockRequest { + mutations, + ..Default::default() + }, + required_version, + ); + } + } + + #[test] + fn transaction_metadata_is_reselected_per_store() { + let mut request = kvrpcpb::GetRequest::default(); + request.context = Some(kvrpcpb::Context { + txn_protocol_version: 99, + request_origin: kvrpcpb::RequestOrigin::Unknown as i32, + ..Default::default() + }); + + for (max, origin, expected_version) in [ + (0, kvrpcpb::RequestOrigin::TiFlash, 0), + (2, kvrpcpb::RequestOrigin::Br, 1), + ] { + request + .prepare_txn_rpc(TxnProtocolVersionRange { min: 0, max }, 1, origin as i32) + .unwrap(); + let context = request.context.as_ref().unwrap(); + assert_eq!(context.txn_protocol_version, expected_version); + assert_eq!( + context.request_origin, + kvrpcpb::RequestOrigin::TiFlash as i32 + ); + } + } + + #[rstest::rstest] + #[case::client_below_required(0, 2, kvrpcpb::Op::SharedLock)] + #[case::store_below_required(0, 1, kvrpcpb::Op::SharedLock)] + #[case::client_below_store_min(2, 2, kvrpcpb::Op::Put)] + #[case::malformed_range(2, 1, kvrpcpb::Op::Put)] + fn incompatible_selection_does_not_modify_context( + #[case] min: u32, + #[case] max: u32, + #[case] op: kvrpcpb::Op, + #[values(false, true)] existing_context: bool, + ) { + let context = existing_context.then_some(kvrpcpb::Context { + txn_protocol_version: 99, + request_origin: kvrpcpb::RequestOrigin::TiFlash as i32, + ..Default::default() + }); + let mut request = kvrpcpb::PrewriteRequest { + context: context.clone(), + mutations: vec![kvrpcpb::Mutation { + op: op as i32, + ..Default::default() + }], + ..Default::default() + }; + let error = Request::prepare_txn_rpc( + &mut request, + TxnProtocolVersionRange { min, max }, + 1, + kvrpcpb::RequestOrigin::Br as i32, + ) + .unwrap_err(); + assert!(matches!(error, Error::IncompatibleRequest(incompatible) + if incompatible.reason == crate::proto::errorpb::IncompatibleRequestReason::Unknown as i32)); + assert_eq!(request.context, context); + } +} diff --git a/src/transaction/client.rs b/src/transaction/client.rs index 62f24859..9d64950d 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -15,11 +15,11 @@ use crate::request::EncodeKeyspace; use crate::request::KeyMode; use crate::request::Keyspace; use crate::request::Plan; +use crate::request::PlanContext; use crate::timestamp::TimestampExt; use crate::transaction::lock::ResolveLocksOptions; use crate::transaction::lowering::new_scan_lock_request; use crate::transaction::lowering::new_unsafe_destroy_range_request; -use crate::transaction::resolve_locks; use crate::transaction::ResolveLocksContext; use crate::transaction::Snapshot; use crate::transaction::Transaction; @@ -56,14 +56,14 @@ const SCAN_LOCK_BATCH_SIZE: u32 = 1024; /// awaited to execute. pub struct Client { pd: Arc, - keyspace: Keyspace, + plan_context: PlanContext, } impl Clone for Client { fn clone(&self) -> Self { Self { pd: self.pd.clone(), - keyspace: self.keyspace, + plan_context: self.plan_context, } } } @@ -125,7 +125,12 @@ impl Client { } None => Keyspace::Disable, }; - Ok(Client { pd, keyspace }) + Ok(Client { + pd, + plan_context: PlanContext::new(keyspace) + .with_request_origin(config.request_origin) + .with_default_txn_protocol_version(config.default_txn_protocol_version), + }) } /// Create a transactional [`Client`] that uses API V2 without adding or removing any API V2 @@ -147,7 +152,9 @@ impl Client { let pd = Arc::new(PdRpcClient::connect(&pd_endpoints, config.clone(), true).await?); Ok(Client { pd, - keyspace: Keyspace::ApiV2NoPrefix, + plan_context: PlanContext::new(Keyspace::ApiV2NoPrefix) + .with_request_origin(config.request_origin) + .with_default_txn_protocol_version(config.default_txn_protocol_version), }) } @@ -299,14 +306,17 @@ impl Client { // scan all locks with ts <= safepoint let ctx = ResolveLocksContext::default(); let backoff = Backoff::equal_jitter_backoff(100, 10000, 50); - let range = range.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let range = range + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); let req = new_scan_lock_request(range, safepoint, options.batch_size); - let plan = crate::request::PlanBuilder::new(self.pd.clone(), self.keyspace, req) - .cleanup_locks(ctx.clone(), options, backoff, self.keyspace) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .extract_error() - .merge(crate::request::Collect) - .plan(); + let plan = + crate::request::PlanBuilder::new_with_context(self.pd.clone(), req, self.plan_context) + .cleanup_locks(ctx.clone(), options, backoff) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .extract_error() + .merge(crate::request::Collect) + .plan(); plan.execute().await } @@ -319,13 +329,19 @@ impl Client { ) -> Result> { use crate::request::TruncateKeyspace; - let range = range.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let range = range + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); let req = new_scan_lock_request(range, safepoint, batch_size); - let plan = crate::request::PlanBuilder::new(self.pd.clone(), self.keyspace, req) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .merge(crate::request::Collect) - .plan(); - Ok(plan.execute().await?.truncate_keyspace(self.keyspace)) + let plan = + crate::request::PlanBuilder::new_with_context(self.pd.clone(), req, self.plan_context) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .merge(crate::request::Collect) + .plan(); + Ok(plan + .execute() + .await? + .truncate_keyspace(self.plan_context.keyspace())) } /// Resolves the given locks and returns any that remain live. @@ -343,14 +359,14 @@ impl Client { let mut live_locks = locks; loop { - let resolved_locks = resolve_locks( - live_locks.encode_keyspace(self.keyspace, KeyMode::Txn), + let resolved_locks = crate::transaction::lock::resolve_locks_with_context( + live_locks.encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn), timestamp.clone(), self.pd.clone(), - self.keyspace, + self.plan_context, ) .await?; - live_locks = resolved_locks.truncate_keyspace(self.keyspace); + live_locks = resolved_locks.truncate_keyspace(self.plan_context.keyspace()); if live_locks.is_empty() { return Ok(live_locks); } @@ -372,16 +388,19 @@ impl Client { /// /// This interface is intended for special scenarios that resemble operations like "drop table" or "drop database" in TiDB. pub async fn unsafe_destroy_range(&self, range: impl Into) -> Result<()> { - let range = range.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let range = range + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); let req = new_unsafe_destroy_range_request(range); - let plan = crate::request::PlanBuilder::new(self.pd.clone(), self.keyspace, req) - .all_stores(DEFAULT_STORE_BACKOFF) - .merge(crate::request::Collect) - .plan(); + let plan = + crate::request::PlanBuilder::new_with_context(self.pd.clone(), req, self.plan_context) + .all_stores(DEFAULT_STORE_BACKOFF) + .merge(crate::request::Collect) + .plan(); plan.execute().await } fn new_transaction(&self, timestamp: Timestamp, options: TransactionOptions) -> Transaction { - Transaction::new(timestamp, self.pd.clone(), options, self.keyspace) + Transaction::new_with_context(timestamp, self.pd.clone(), options, self.plan_context) } } diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index 98bb3683..7b90d7a7 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -17,16 +17,18 @@ use crate::backoff::OPTIMISTIC_BACKOFF; use crate::kv::HexRepr; use crate::pd::PdClient; +use crate::common::is_grpc_error; use crate::proto::kvrpcpb; use crate::proto::kvrpcpb::TxnInfo; use crate::proto::pdpb::Timestamp; use crate::region::RegionVerId; use crate::request::plan::handle_region_error; -use crate::request::plan::is_grpc_error; use crate::request::Collect; use crate::request::CollectSingle; use crate::request::Keyspace; use crate::request::Plan; +use crate::request::PlanContext; +use crate::store::reload_txn_protocol_range; use crate::store::RegionStore; use crate::timestamp::TimestampExt; use crate::transaction::requests; @@ -78,12 +80,13 @@ pub(crate) fn reject_shared_locks(locks: &[kvrpcpb::LockInfo]) -> Result<()> { /// which means the key is finally either committed or rolled back, before we read the value of /// the key. We first use `CheckTxnStatus` to get the transaction's final status (committed or /// rolled back), then use `ResolveLock` to resolve the remaining locks in the transaction. -pub async fn resolve_locks( +pub(crate) async fn resolve_locks_with_context( locks: Vec, timestamp: Timestamp, pd_client: Arc, - keyspace: Keyspace, + plan_context: PlanContext, ) -> Result /* live_locks */> { + let keyspace = plan_context.keyspace(); debug!("resolving locks"); reject_shared_locks(&locks)?; let ts = pd_client.clone().get_timestamp().await?; @@ -91,7 +94,8 @@ pub async fn resolve_locks( let current_ts = ts.version(); let mut live_locks = Vec::new(); - let mut lock_resolver = LockResolver::new(ResolveLocksContext::default()); + let mut lock_resolver = + LockResolver::new_with_context(ResolveLocksContext::default(), plan_context); // records the commit version of each primary lock (representing the status of the transaction) let mut commit_versions: HashMap = HashMap::new(); @@ -160,7 +164,7 @@ pub async fn resolve_locks( commit_version, lock.is_txn_file, pd_client.clone(), - keyspace, + plan_context, OPTIMISTIC_BACKOFF, ) .await?; @@ -179,7 +183,7 @@ async fn resolve_lock_with_retry( commit_version: u64, is_txn_file: bool, pd_client: Arc, - keyspace: Keyspace, + plan_context: PlanContext, mut backoff: Backoff, ) -> Result { debug!("resolving locks with retry"); @@ -191,24 +195,28 @@ async fn resolve_lock_with_retry( let ver_id = store.region_with_leader.ver_id(); let request = requests::new_resolve_lock_request(start_version, commit_version, is_txn_file); - let plan_builder = - match crate::request::PlanBuilder::new(pd_client.clone(), keyspace, request) - .single_region_with_store(store.clone()) - .await - { - Ok(plan_builder) => plan_builder, - Err(Error::LeaderNotFound { region }) => { - pd_client.invalidate_region_cache(region.clone()).await; - match backoff.next_delay_duration() { - Some(duration) => { - sleep(duration).await; - continue; - } - None => return Err(Error::LeaderNotFound { region }), + let requirement = crate::store::Request::txn_protocol_requirement(&request); + let plan_builder = match crate::request::PlanBuilder::new_with_context( + pd_client.clone(), + request, + plan_context, + ) + .single_region_with_store(store.clone()) + .await + { + Ok(plan_builder) => plan_builder, + Err(Error::LeaderNotFound { region }) => { + pd_client.invalidate_region_cache(region.clone()).await; + match backoff.next_delay_duration() { + Some(duration) => { + sleep(duration).await; + continue; } + None => return Err(Error::LeaderNotFound { region }), } - Err(err) => return Err(err), - }; + } + Err(err) => return Err(err), + }; let plan = plan_builder.extract_error().plan(); match plan.execute().await { Ok(_) => { @@ -256,6 +264,23 @@ async fn resolve_lock_with_retry( } None => return Err(e), }, + Err(Error::IncompatibleRequest(incompatible)) => { + if reload_txn_protocol_range( + pd_client.clone(), + store.store_id, + store.txn_protocol_version_range, + plan_context.default_txn_protocol_version(), + requirement, + &incompatible, + &mut backoff, + ) + .await + .is_some() + { + continue; + } + return Err(Error::IncompatibleRequest(incompatible)); + } Err(e) => return Err(e), } } @@ -313,11 +338,22 @@ impl ResolveLocksContext { pub struct LockResolver { ctx: ResolveLocksContext, + plan_context: PlanContext, } impl LockResolver { pub fn new(ctx: ResolveLocksContext) -> Self { - Self { ctx } + Self::new_with_context(ctx, PlanContext::new(Keyspace::Disable)) + } + + pub(crate) fn new_with_context(ctx: ResolveLocksContext, plan_context: PlanContext) -> Self { + Self { ctx, plan_context } + } + + /// The public resolver API accepts a keyspace per call. Keep its RPC context + /// in sync with that keyspace at the single conversion point. + fn context_for(&self, keyspace: Keyspace) -> PlanContext { + self.plan_context.with_keyspace(keyspace) } /// _Cleanup_ the given locks. Returns whether all the given locks are resolved. @@ -492,12 +528,16 @@ impl LockResolver { resolving_pessimistic_lock, is_txn_file, ); - let plan = crate::request::PlanBuilder::new(pd_client.clone(), keyspace, req) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .merge(CollectSingle) - .extract_error() - .post_process_default() - .plan(); + let plan = crate::request::PlanBuilder::new_with_context( + pd_client.clone(), + req, + self.context_for(keyspace), + ) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .merge(CollectSingle) + .extract_error() + .post_process_default() + .plan(); let mut status: TransactionStatus = match plan.execute().await { Ok(status) => status, Err(Error::ExtractedErrors(mut errors)) => match errors.pop() { @@ -531,11 +571,15 @@ impl LockResolver { txn_id: u64, ) -> Result { let req = new_check_secondary_locks_request(keys, txn_id); - let plan = crate::request::PlanBuilder::new(pd_client.clone(), keyspace, req) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .extract_error() - .merge(Collect) - .plan(); + let plan = crate::request::PlanBuilder::new_with_context( + pd_client.clone(), + req, + self.context_for(keyspace), + ) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .extract_error() + .merge(Collect) + .plan(); plan.execute().await } @@ -543,18 +587,46 @@ impl LockResolver { &mut self, pd_client: Arc, keyspace: Keyspace, - store: RegionStore, + mut store: RegionStore, txn_infos: Vec, ) -> Result { let ver_id = store.region_with_leader.ver_id(); - let request = requests::new_batch_resolve_lock_request(txn_infos.clone()); - let plan = crate::request::PlanBuilder::new(pd_client.clone(), keyspace, request) + let mut backoff = DEFAULT_REGION_BACKOFF; + let plan_context = self.context_for(keyspace); + loop { + let request = requests::new_batch_resolve_lock_request(txn_infos.clone()); + let requirement = crate::store::Request::txn_protocol_requirement(&request); + let plan = crate::request::PlanBuilder::new_with_context( + pd_client.clone(), + request, + plan_context, + ) .single_region_with_store(store.clone()) .await? .extract_error() .plan(); - let _ = plan.execute().await?; - Ok(ver_id) + match plan.execute().await { + Ok(_) => return Ok(ver_id), + Err(Error::IncompatibleRequest(incompatible)) => { + if let Some(range) = reload_txn_protocol_range( + pd_client.clone(), + store.store_id, + store.txn_protocol_version_range, + plan_context.default_txn_protocol_version(), + requirement, + &incompatible, + &mut backoff, + ) + .await + { + store.txn_protocol_version_range = range; + continue; + } + return Err(Error::IncompatibleRequest(incompatible)); + } + Err(error) => return Err(error), + } + } } #[allow(clippy::too_many_arguments)] @@ -652,6 +724,77 @@ mod tests { use crate::mock::MockPdClient; use crate::proto::errorpb; + #[rstest::rstest] + #[case(crate::mock::ReloadOutcome::Changed)] + #[case(crate::mock::ReloadOutcome::Unchanged)] + #[tokio::test] + async fn lock_admission_reload_reselects_before_resending( + #[values(false, true)] batch: bool, + #[case] outcome: crate::mock::ReloadOutcome, + ) { + let versions = Arc::new(std::sync::Mutex::new(Vec::new())); + let recorded = versions.clone(); + let incompatible = crate::mock::upper_admission_rejection(); + let rejection = incompatible.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + let request = request + .downcast_ref::() + .unwrap(); + let version = request.context.as_ref().unwrap().txn_protocol_version; + recorded.lock().unwrap().push(version); + Ok(Box::new(kvrpcpb::ResolveLockResponse { + region_error: (version == 1).then_some(errorpb::Error { + incompatible_request: Some(rejection.clone()), + ..Default::default() + }), + ..Default::default() + }) as Box) + }); + let fixture = crate::mock::ProtocolTestClient::new(client, outcome); + let context = PlanContext::new(Keyspace::Disable); + let result = if batch { + let store = fixture + .pd_client + .clone() + .store_for_key(&vec![1].into()) + .await + .unwrap(); + LockResolver::new_with_context(ResolveLocksContext::default(), context) + .batch_resolve_locks( + fixture.pd_client, + Keyspace::Disable, + store, + vec![TxnInfo { + txn: 1, + status: 2, + ..Default::default() + }], + ) + .await + } else { + resolve_lock_with_retry( + &vec![1], + 1, + 2, + false, + fixture.pd_client, + context, + Backoff::no_jitter_backoff(0, 0, 1), + ) + .await + }; + if matches!(outcome, crate::mock::ReloadOutcome::Changed) { + assert!(result.is_ok()); + assert_eq!(*versions.lock().unwrap(), vec![1, 0]); + } else { + assert!( + matches!(result, Err(Error::IncompatibleRequest(error)) if *error == incompatible) + ); + assert_eq!(*versions.lock().unwrap(), vec![1]); + } + assert_eq!(*fixture.reloads.lock().unwrap(), vec![41]); + } + #[test] fn shared_locks_are_refused_never_misresolved() { let plain = kvrpcpb::LockInfo { @@ -681,6 +824,45 @@ mod tests { assert!(reject_shared_locks(&[by_op]).is_err()); } + #[tokio::test] + async fn public_lock_resolver_uses_the_passed_keyspace() { + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + |req: &dyn Any| { + let req = req + .downcast_ref::() + .expect("expected CheckTxnStatusRequest"); + assert_eq!( + req.context.as_ref().unwrap().api_version, + kvrpcpb::ApiVersion::V2 as i32 + ); + Ok(Box::new(kvrpcpb::CheckTxnStatusResponse { + commit_version: 2, + action: kvrpcpb::Action::NoAction as i32, + ..Default::default() + }) as Box) + }, + ))); + let mut resolver = LockResolver::new(ResolveLocksContext::default()); + + let status = resolver + .check_txn_status( + client, + Keyspace::Enable { keyspace_id: 7 }, + 1, + vec![1], + 0, + 0, + false, + false, + false, + false, + ) + .await + .unwrap(); + + assert!(matches!(&status.kind, TransactionStatusKind::Committed(_))); + } + #[rstest::rstest] #[case(Keyspace::Disable)] #[case(Keyspace::Enable { keyspace_id: 0 })] @@ -714,10 +896,17 @@ mod tests { let key = vec![1]; let region1 = MockPdClient::region1(); - let resolved_region = - resolve_lock_with_retry(&key, 1, 2, false, client.clone(), keyspace, backoff.clone()) - .await - .unwrap(); + let resolved_region = resolve_lock_with_retry( + &key, + 1, + 2, + false, + client.clone(), + PlanContext::new(keyspace), + backoff.clone(), + ) + .await + .unwrap(); assert_eq!(region1.ver_id(), resolved_region); // Test resolve lock over retry limit @@ -727,9 +916,17 @@ mod tests { ) .unwrap(); let key = vec![100]; - resolve_lock_with_retry(&key, 3, 4, false, client, keyspace, backoff) - .await - .expect_err("should return error"); + resolve_lock_with_retry( + &key, + 3, + 4, + false, + client, + PlanContext::new(keyspace), + backoff, + ) + .await + .expect_err("should return error"); } #[tokio::test] @@ -765,9 +962,14 @@ mod tests { lock.lock_version = 1; lock.lock_ttl = 100; // not expired under MockPdClient's Timestamp::default() - let live_locks = resolve_locks(vec![lock], Timestamp::default(), client, Keyspace::Disable) - .await - .unwrap(); + let live_locks = resolve_locks_with_context( + vec![lock], + Timestamp::default(), + client, + PlanContext::new(Keyspace::Disable), + ) + .await + .unwrap(); assert!(live_locks.is_empty()); assert_eq!(check_txn_status_count.load(Ordering::SeqCst), 1); diff --git a/src/transaction/mod.rs b/src/transaction/mod.rs index 86612a24..f8f037a1 100644 --- a/src/transaction/mod.rs +++ b/src/transaction/mod.rs @@ -10,7 +10,7 @@ pub use client::Client; pub use client::ProtoLockInfo; -pub(crate) use lock::resolve_locks; +pub(crate) use lock::resolve_locks_with_context; pub(crate) use lock::HasLocks; pub use snapshot::Snapshot; pub use sync_client::SyncTransactionClient; diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs index db524f16..36e2bb5c 100644 --- a/src/transaction/requests.rs +++ b/src/transaction/requests.rs @@ -46,6 +46,7 @@ use crate::timestamp::TimestampExt; use crate::transaction::requests::kvrpcpb::prewrite_request::PessimisticAction; use crate::transaction::HasLocks; use crate::util::iter::FlatMapOkIterExt; +use crate::Error; use crate::KvPair; use crate::Result; use crate::Value; @@ -457,9 +458,19 @@ impl Merge, ) -> Result { if input.iter().any(Result::is_err) { - let (success, mut errors): (Vec<_>, Vec<_>) = - input.into_iter().partition(Result::is_ok); - let first_err = errors.pop().unwrap(); + let (success, errors): (Vec<_>, Vec<_>) = input.into_iter().partition(Result::is_ok); + let mut errors = errors + .into_iter() + .map(Result::unwrap_err) + .collect::>(); + let selected_error = errors + .iter() + .enumerate() + .filter(|(_, error)| error.priority() != crate::common::ErrorPriority::Ordinary) + .min_by_key(|(_, error)| error.priority()) + .map(|(index, _)| index) + .unwrap_or(errors.len() - 1); + let error = errors.swap_remove(selected_error); let success_keys = success .into_iter() .map(Result::unwrap) @@ -468,7 +479,7 @@ impl Merge Result<()> { + Ok(()) + } } impl HasLocks for kvrpcpb::UnsafeDestroyRangeResponse {} @@ -883,8 +896,20 @@ impl Merge for Collect { type Out = (); fn merge(&self, input: Vec>) -> Result { - let _: Vec = - input.into_iter().collect::>>()?; + let mut selected_error: Option = None; + for result in input { + if let Err(error) = result { + if selected_error + .as_ref() + .is_none_or(|selected| error.priority() < selected.priority()) + { + selected_error = Some(error); + } + } + } + if let Some(error) = selected_error { + return Err(error); + } Ok(()) } } @@ -899,6 +924,63 @@ mod tests { use crate::request::ResponseWithShard; use crate::KvPair; + #[rstest::rstest] + #[case(crate::common::ErrorPriority::Undetermined)] + #[case(crate::common::ErrorPriority::Incompatible)] + #[case(crate::common::ErrorPriority::Ordinary)] + fn merges_preserve_priority_and_pessimistic_success_keys( + #[case] priority: crate::common::ErrorPriority, + ) { + use crate::common::ErrorPriority; + use crate::Error; + let errors = || { + let mut errors = vec![Error::StringError("first ordinary".into())]; + if priority != ErrorPriority::Ordinary { + errors.push(Error::IncompatibleRequest(Box::new( + crate::mock::upper_admission_rejection(), + ))); + } + if priority == ErrorPriority::Undetermined { + errors.push(Error::from(crate::proto::errorpb::Error { + undetermined_result: Some(Default::default()), + ..Default::default() + })); + } + errors.push(Error::StringError("last ordinary".into())); + errors + }; + let mut input = errors().into_iter().map(Err).collect::>(); + input.push(Ok(ResponseWithShard( + kvrpcpb::PessimisticLockResponse::default(), + vec![kvrpcpb::Mutation { + key: b"cleanup-me".to_vec(), + ..Default::default() + }], + ))); + match CollectWithShard.merge(input).unwrap_err() { + PessimisticLockError { + inner, + success_keys, + } => { + assert_eq!(inner.priority(), priority); + assert_eq!(success_keys, vec![b"cleanup-me".to_vec()]); + if priority == ErrorPriority::Ordinary { + assert!( + matches!(*inner, Error::StringError(message) if message == "last ordinary") + ); + } + } + other => panic!("wrong wrapper: {other:?}"), + } + let input: Vec> = + errors().into_iter().map(Err).collect(); + let error = crate::request::Collect.merge(input).unwrap_err(); + assert_eq!(error.priority(), priority); + if priority == ErrorPriority::Ordinary { + assert!(matches!(error, Error::StringError(message) if message == "first ordinary")); + } + } + #[tokio::test] async fn test_merge_pessimistic_lock_response() { let (key1, key2, key3, key4) = (b"key1", b"key2", b"key3", b"key4"); diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 7f57146a..2c98c058 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -14,6 +14,7 @@ use tokio::time::Duration; use crate::backoff::Backoff; use crate::backoff::DEFAULT_REGION_BACKOFF; +use crate::common::ErrorPriority; use crate::pd::PdClient; use crate::pd::PdRpcClient; use crate::proto::kvrpcpb; @@ -24,9 +25,11 @@ use crate::request::CollectSingle; use crate::request::CollectWithShard; use crate::request::EncodeKeyspace; use crate::request::KeyMode; +#[cfg(test)] use crate::request::Keyspace; use crate::request::Plan; use crate::request::PlanBuilder; +use crate::request::PlanContext; use crate::request::RetryOptions; use crate::request::TruncateKeyspace; use crate::timestamp::TimestampExt; @@ -87,7 +90,7 @@ pub struct Transaction { buffer: Buffer, rpc: Arc, options: TransactionOptions, - keyspace: Keyspace, + plan_context: PlanContext, is_heartbeat_started: bool, /// Set once the transaction enters the commit path (`StartedCommit`), where /// prewrite may place 2PC locks. Kept as a dedicated flag because the status @@ -98,11 +101,21 @@ pub struct Transaction { } impl Transaction { + #[cfg(test)] pub(crate) fn new( timestamp: Timestamp, rpc: Arc, options: TransactionOptions, keyspace: Keyspace, + ) -> Transaction { + Self::new_with_context(timestamp, rpc, options, PlanContext::new(keyspace)) + } + + pub(crate) fn new_with_context( + timestamp: Timestamp, + rpc: Arc, + options: TransactionOptions, + plan_context: PlanContext, ) -> Transaction { let status = if options.read_only { TransactionStatus::ReadOnly @@ -115,7 +128,7 @@ impl Transaction { buffer: Buffer::new(options.is_pessimistic()), rpc, options, - keyspace, + plan_context, is_heartbeat_started: false, prewritten: false, start_instant: std::time::Instant::now(), @@ -145,15 +158,17 @@ impl Transaction { self.check_allow_operation().await?; let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); - let key = key.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let key = key + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); let retry_options = self.options.retry_options.clone(); - let keyspace = self.keyspace; + let plan_context = self.plan_context; self.buffer .get_or_else(key, |key| async move { let request = new_get_request(key, timestamp.clone()); - let plan = PlanBuilder::new(rpc, keyspace, request) - .resolve_lock(timestamp, retry_options.lock_backoff, keyspace) + let plan = PlanBuilder::new_with_context(rpc, request, plan_context) + .resolve_lock(timestamp, retry_options.lock_backoff) .retry_multi_region(DEFAULT_REGION_BACKOFF) .merge(CollectSingle) .post_process_default() @@ -213,7 +228,9 @@ impl Transaction { self.lock_keys(iter::once(key.clone())).await?; self.get(key).await } else { - let key = key.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let key = key + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); let mut pairs = self.pessimistic_lock(iter::once(key), true).await?; debug_assert!(pairs.len() <= 1); match pairs.pop() { @@ -278,7 +295,8 @@ impl Transaction { self.check_allow_operation().await?; let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); - let keyspace = self.keyspace; + let keyspace = self.plan_context.keyspace(); + let plan_context = self.plan_context; let keys = keys .into_iter() .map(move |k| k.into().encode_keyspace(keyspace, KeyMode::Txn)); @@ -287,8 +305,8 @@ impl Transaction { self.buffer .batch_get_or_else(keys, move |keys| async move { let request = new_batch_get_request(keys, timestamp.clone()); - let plan = PlanBuilder::new(rpc, keyspace, request) - .resolve_lock(timestamp, retry_options.lock_backoff, keyspace) + let plan = PlanBuilder::new_with_context(rpc, request, plan_context) + .resolve_lock(timestamp, retry_options.lock_backoff) .retry_multi_region(retry_options.region_backoff) .merge(Collect) .plan(); @@ -336,7 +354,7 @@ impl Transaction { self.lock_keys(keys.clone()).await?; Ok(self.batch_get(keys).await?.collect()) } else { - let keyspace = self.keyspace; + let keyspace = self.plan_context.keyspace(); let keys = keys .into_iter() .map(move |k| k.into().encode_keyspace(keyspace, KeyMode::Txn)); @@ -470,7 +488,9 @@ impl Transaction { pub async fn put(&mut self, key: impl Into, value: impl Into) -> Result<()> { trace!("invoking transactional put request"); self.check_allow_operation().await?; - let key = key.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let key = key + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); if self.is_pessimistic() { self.pessimistic_lock(iter::once(key.clone()), false) .await?; @@ -501,7 +521,9 @@ impl Transaction { pub async fn insert(&mut self, key: impl Into, value: impl Into) -> Result<()> { debug!("invoking transactional insert request"); self.check_allow_operation().await?; - let key = key.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let key = key + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); if self.buffer.get(&key).is_some() { return Err(Error::DuplicateKeyInsertion); } @@ -536,7 +558,9 @@ impl Transaction { pub async fn delete(&mut self, key: impl Into) -> Result<()> { debug!("invoking transactional delete request"); self.check_allow_operation().await?; - let key = key.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let key = key + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); if self.is_pessimistic() { self.pessimistic_lock(iter::once(key.clone()), false) .await?; @@ -573,7 +597,7 @@ impl Transaction { self.check_allow_operation().await?; let mutations: Vec = mutations .into_iter() - .map(|mutation| mutation.encode_keyspace(self.keyspace, KeyMode::Txn)) + .map(|mutation| mutation.encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn)) .collect(); if self.is_pessimistic() { self.pessimistic_lock(mutations.iter().map(|m| m.key().clone()), false) @@ -618,7 +642,7 @@ impl Transaction { ) -> Result<()> { debug!("invoking transactional lock_keys request"); self.check_allow_operation().await?; - let keyspace = self.keyspace; + let keyspace = self.plan_context.keyspace(); let keys = keys .into_iter() .map(move |k| k.into().encode_keyspace(keyspace, KeyMode::Txn)); @@ -686,7 +710,7 @@ impl Transaction { self.timestamp.clone(), self.rpc.clone(), self.options.clone(), - self.keyspace, + self.plan_context, self.buffer.get_write_size() as u64, self.start_instant, ) @@ -755,7 +779,7 @@ impl Transaction { self.timestamp.clone(), self.rpc.clone(), self.options.clone(), - self.keyspace, + self.plan_context, self.buffer.get_write_size() as u64, self.start_instant, ) @@ -793,11 +817,10 @@ impl Transaction { primary_key, self.start_instant.elapsed().as_millis() as u64 + MAX_TTL, ); - let plan = PlanBuilder::new(self.rpc.clone(), self.keyspace, request) + let plan = PlanBuilder::new_with_context(self.rpc.clone(), request, self.plan_context) .resolve_lock( self.timestamp.clone(), self.options.retry_options.lock_backoff.clone(), - self.keyspace, ) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .extract_error() @@ -818,8 +841,11 @@ impl Transaction { let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); let retry_options = self.options.retry_options.clone(); - let keyspace = self.keyspace; - let range = range.into().encode_keyspace(self.keyspace, KeyMode::Txn); + let keyspace = self.plan_context.keyspace(); + let plan_context = self.plan_context; + let range = range + .into() + .encode_keyspace(self.plan_context.keyspace(), KeyMode::Txn); self.buffer .scan_and_fetch( @@ -835,8 +861,8 @@ impl Transaction { key_only, reverse, ); - let plan = PlanBuilder::new(rpc, keyspace, request) - .resolve_lock(timestamp, retry_options.lock_backoff, keyspace) + let plan = PlanBuilder::new_with_context(rpc, request, plan_context) + .resolve_lock(timestamp, retry_options.lock_backoff) .retry_multi_region(retry_options.region_backoff) .merge(Collect) .plan(); @@ -894,11 +920,10 @@ impl Transaction { for_update_ts.clone(), need_value, ); - let plan = PlanBuilder::new(self.rpc.clone(), self.keyspace, request) + let plan = PlanBuilder::new_with_context(self.rpc.clone(), request, self.plan_context) .resolve_lock( self.timestamp.clone(), self.options.retry_options.lock_backoff.clone(), - self.keyspace, ) .preserve_shard() .retry_multi_region_preserve_results(self.options.retry_options.region_backoff.clone()) @@ -918,10 +943,30 @@ impl Transaction { self.timestamp.version(), for_update_ts.version(), ); + let key_count = success_keys.len(); let keys = success_keys.into_iter().map(Key::from); - self.pessimistic_lock_rollback(keys, self.timestamp.clone(), for_update_ts) - .await?; - Err(*inner) + match self + .pessimistic_lock_rollback(keys, self.timestamp.clone(), for_update_ts) + .await + { + Ok(()) => Err(*inner), + Err(rollback_error) => { + warn!( + "failed to roll back partially-acquired pessimistic locks, start_ts: {}, keys: {}, lock error: {}, rollback error: {}", + self.timestamp.version(), key_count, inner, rollback_error, + ); + // Preserve high-priority lock errors unless cleanup reports + // a higher-priority error. Ordinary errors retain the existing + // behavior of returning the cleanup failure. + if inner.priority() != ErrorPriority::Ordinary + && inner.priority() <= rollback_error.priority() + { + Err(*inner) + } else { + Err(rollback_error) + } + } + } } _ => Err(err), } @@ -962,11 +1007,10 @@ impl Transaction { start_version.clone(), for_update_ts, ); - let plan = PlanBuilder::new(self.rpc.clone(), self.keyspace, req) + let plan = PlanBuilder::new_with_context(self.rpc.clone(), req, self.plan_context) .resolve_lock( start_version, self.options.retry_options.lock_backoff.clone(), - self.keyspace, ) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .extract_error() @@ -1014,7 +1058,7 @@ impl Transaction { HeartbeatOption::FixedTime(heartbeat_interval) => heartbeat_interval, }; let start_instant = self.start_instant; - let keyspace = self.keyspace; + let plan_context = self.plan_context; debug!( "starting auto-heartbeat, start_ts: {}, interval: {:?}", self.timestamp.version(), @@ -1040,7 +1084,7 @@ impl Transaction { primary_key.clone(), start_instant.elapsed().as_millis() as u64 + MAX_TTL, ); - let plan = PlanBuilder::new(rpc.clone(), keyspace, request) + let plan = PlanBuilder::new_with_context(rpc.clone(), request, plan_context) .retry_multi_region(region_backoff.clone()) .merge(CollectSingle) .plan(); @@ -1329,7 +1373,7 @@ struct Committer { start_version: Timestamp, rpc: Arc, options: TransactionOptions, - keyspace: Keyspace, + plan_context: PlanContext, #[new(default)] undetermined: bool, write_size: u64, @@ -1441,11 +1485,11 @@ impl Committer { .collect(); // FIXME set max_commit_ts and min_commit_ts - let builder = PlanBuilder::new(self.rpc.clone(), self.keyspace, request).resolve_lock( - self.start_version.clone(), - self.options.retry_options.lock_backoff.clone(), - self.keyspace, - ); + let builder = PlanBuilder::new_with_context(self.rpc.clone(), request, self.plan_context) + .resolve_lock( + self.start_version.clone(), + self.options.retry_options.lock_backoff.clone(), + ); // With async commit or 1PC this prewrite IS the commit point, so an unknown // apply outcome must not be retried away or overwritten by a later error // (stricter than client-go, in the safe direction; commit() classifies it). @@ -1495,11 +1539,10 @@ impl Committer { self.start_version.clone(), commit_version.clone(), ); - let plan = PlanBuilder::new(self.rpc.clone(), self.keyspace, req) + let plan = PlanBuilder::new_with_context(self.rpc.clone(), req, self.plan_context) .resolve_lock( self.start_version.clone(), self.options.retry_options.lock_backoff.clone(), - self.keyspace, ) // The primary commit is THE commit point: client-go returns // ErrResultUndetermined here rather than retrying (commit.go), so an @@ -1634,12 +1677,8 @@ impl Committer { .filter(|key| &primary_key != key); new_commit_request(keys, start_version.clone(), commit_version) }; - let plan = PlanBuilder::new(self.rpc, self.keyspace, req) - .resolve_lock( - start_version, - self.options.retry_options.lock_backoff, - self.keyspace, - ) + let plan = PlanBuilder::new_with_context(self.rpc, req, self.plan_context) + .resolve_lock(start_version, self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() .plan(); @@ -1681,13 +1720,13 @@ impl Committer { let lock_backoff = self.options.retry_options.lock_backoff.clone(); let region_backoff = self.options.retry_options.region_backoff.clone(); let rpc = self.rpc; - let keyspace = self.keyspace; + let plan_context = self.plan_context; match self.options.kind { TransactionKind::Pessimistic(for_update_ts) if !prewritten => { let req = new_pessimistic_rollback_request(keys, start_version.clone(), for_update_ts); - let plan = PlanBuilder::new(rpc, keyspace, req) - .resolve_lock(start_version, lock_backoff, keyspace) + let plan = PlanBuilder::new_with_context(rpc, req, plan_context) + .resolve_lock(start_version, lock_backoff) .retry_multi_region(region_backoff) .extract_error() .plan(); @@ -1697,8 +1736,8 @@ impl Committer { // both pessimistic and 2PC locks by start_ts. _ => { let req = new_batch_rollback_request(keys, start_version.clone()); - let plan = PlanBuilder::new(rpc, keyspace, req) - .resolve_lock(start_version, lock_backoff, keyspace) + let plan = PlanBuilder::new_with_context(rpc, req, plan_context) + .resolve_lock(start_version, lock_backoff) .retry_multi_region(region_backoff) .extract_error() .plan(); @@ -1729,7 +1768,7 @@ impl Committer { /// dispatch failure takes since the tonic migration; `Error::Grpc` is only connection /// establishment. fn is_commit_outcome_unknown(e: &Error) -> bool { - crate::request::plan::is_grpc_error(e) || crate::request::plan::is_undetermined_region_error(e) + crate::common::is_grpc_error(e) || crate::common::is_undetermined_region_error(e) } #[derive(Debug, PartialEq, Eq, Clone, Copy)] @@ -1779,6 +1818,7 @@ mod tests { use fail::FailScenario; + use crate::common::ErrorPriority; use crate::mock::MockKvClient; use crate::mock::MockPdClient; use crate::proto::kvrpcpb; @@ -1789,6 +1829,218 @@ mod tests { use crate::Transaction; use crate::TransactionOptions; + #[derive(Clone, Copy, Debug)] + enum LockRollbackOutcome { + Success, + Grpc, + Incompatible, + Undetermined, + } + + #[rstest::rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn commit_point_undetermined_bypasses_lock_resolution( + #[case] one_pc: bool, + #[values(false, true)] resolve_enabled: bool, + ) { + let attempts = Arc::new(AtomicUsize::new(0)); + let recorded = attempts.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + let request = request + .downcast_ref::() + .expect("undetermined prewrite must not trigger auxiliary RPCs"); + assert_eq!(request.try_one_pc, one_pc); + assert_eq!(request.use_async_commit, !one_pc); + recorded.fetch_add(1, Ordering::SeqCst); + Ok(Box::new(kvrpcpb::PrewriteResponse { + region_error: Some(crate::proto::errorpb::Error { + message: "original undetermined prewrite".into(), + undetermined_result: Some(Default::default()), + ..Default::default() + }), + errors: vec![kvrpcpb::KeyError { + locked: Some(kvrpcpb::LockInfo { + key: vec![1], + primary_lock: vec![1], + lock_version: 1, + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }) as Box) + }); + let options = TransactionOptions::new_optimistic() + .no_resolve_regions() + .heartbeat_option(HeartbeatOption::NoHeartbeat) + .drop_check(super::CheckLevel::None); + let options = if one_pc { + options.try_one_pc() + } else { + options.use_async_commit() + }; + let options = if resolve_enabled { + options + } else { + options.no_resolve_locks() + }; + let mut txn = Transaction::new( + Timestamp::from_version(123), + Arc::new(MockPdClient::new(client)), + options, + Keyspace::Disable, + ); + txn.put(vec![1], vec![2]).await.unwrap(); + let Err(crate::Error::UndeterminedError(error)) = txn.commit().await else { + panic!("expected an undetermined commit outcome"); + }; + assert!(matches!(*error, crate::Error::RegionError(error) + if error.undetermined_result.is_some() && error.message == "original undetermined prewrite")); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + } + + #[rstest::rstest] + #[case(ErrorPriority::Incompatible, LockRollbackOutcome::Success, true)] + #[case(ErrorPriority::Incompatible, LockRollbackOutcome::Grpc, true)] + #[case(ErrorPriority::Incompatible, LockRollbackOutcome::Incompatible, true)] + #[case(ErrorPriority::Incompatible, LockRollbackOutcome::Undetermined, false)] + #[case(ErrorPriority::Undetermined, LockRollbackOutcome::Success, true)] + #[case(ErrorPriority::Undetermined, LockRollbackOutcome::Grpc, true)] + #[case(ErrorPriority::Undetermined, LockRollbackOutcome::Incompatible, true)] + #[case(ErrorPriority::Undetermined, LockRollbackOutcome::Undetermined, true)] + #[case(ErrorPriority::Ordinary, LockRollbackOutcome::Success, true)] + #[case(ErrorPriority::Ordinary, LockRollbackOutcome::Grpc, false)] + #[case(ErrorPriority::Ordinary, LockRollbackOutcome::Incompatible, false)] + #[case(ErrorPriority::Ordinary, LockRollbackOutcome::Undetermined, false)] + #[tokio::test] + async fn partial_pessimistic_lock_preserves_error_priority_after_rollback( + #[case] original_priority: ErrorPriority, + #[case] rollback: LockRollbackOutcome, + #[case] preserve_original: bool, + ) { + fn region_error(priority: ErrorPriority, message: &str) -> crate::proto::errorpb::Error { + let mut error = crate::proto::errorpb::Error { + message: message.into(), + ..Default::default() + }; + match priority { + ErrorPriority::Undetermined => error.undetermined_result = Some(Default::default()), + ErrorPriority::Incompatible => { + error.incompatible_request = Some(crate::proto::errorpb::IncompatibleRequest { + message: message.into(), + ..Default::default() + }); + } + ErrorPriority::Ordinary => unreachable!(), + } + error + } + + let lock_keys = Arc::new(Mutex::new(Vec::new())); + let locked = lock_keys.clone(); + let rollback_keys = Arc::new(Mutex::new(Vec::new())); + let rolled_back = rollback_keys.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + if let Some(request) = request.downcast_ref::() { + let keys = request + .mutations + .iter() + .map(|m| m.key.clone()) + .collect::>(); + locked.lock().unwrap().extend(keys.clone()); + assert_eq!(keys.len(), 1); + let mut response = kvrpcpb::PessimisticLockResponse::default(); + if keys[0] == vec![11] { + if original_priority == ErrorPriority::Ordinary { + response.errors.push(kvrpcpb::KeyError { + abort: "original lock error".into(), + ..Default::default() + }); + } else { + response.region_error = + Some(region_error(original_priority, "original lock error")); + } + } else { + assert_eq!(keys[0], vec![1]); + } + return Ok(Box::new(response) as Box); + } + let request = request + .downcast_ref::() + .expect("expected a rollback of the successfully locked shard"); + rolled_back.lock().unwrap().push(request.keys.clone()); + let mut response = kvrpcpb::PessimisticRollbackResponse::default(); + match rollback { + LockRollbackOutcome::Success => {} + LockRollbackOutcome::Grpc => { + return Err(crate::Error::GrpcAPI(tonic::Status::unavailable( + "rollback error", + ))); + } + LockRollbackOutcome::Incompatible => { + response.region_error = + Some(region_error(ErrorPriority::Incompatible, "rollback error")); + } + LockRollbackOutcome::Undetermined => { + response.region_error = + Some(region_error(ErrorPriority::Undetermined, "rollback error")); + } + } + Ok(Box::new(response) as Box) + }); + let mut txn = Transaction::new( + Timestamp::from_version(123), + Arc::new(MockPdClient::new(client)), + TransactionOptions::new_pessimistic() + .no_resolve_locks() + .no_resolve_regions() + .heartbeat_option(HeartbeatOption::NoHeartbeat) + .drop_check(super::CheckLevel::None), + Keyspace::Disable, + ); + let error = txn.lock_keys([vec![1], vec![11]]).await.unwrap_err(); + + let expected_message = if preserve_original { + "original lock error" + } else { + "rollback error" + }; + let expected_priority = if preserve_original { + original_priority + } else { + match rollback { + LockRollbackOutcome::Grpc => ErrorPriority::Ordinary, + LockRollbackOutcome::Incompatible => ErrorPriority::Incompatible, + LockRollbackOutcome::Undetermined => ErrorPriority::Undetermined, + LockRollbackOutcome::Success => unreachable!(), + } + }; + assert_eq!(error.priority(), expected_priority); + match error { + crate::Error::IncompatibleRequest(error) => assert_eq!(error.message, expected_message), + crate::Error::RegionError(error) => { + assert!(error.undetermined_result.is_some()); + assert_eq!(error.message, expected_message); + } + crate::Error::GrpcAPI(status) => { + assert_eq!(status.code(), tonic::Code::Unavailable); + assert_eq!(status.message(), expected_message); + } + crate::Error::MultipleKeyErrors(errors) => { + assert_eq!(errors.len(), 1); + assert!(matches!(&errors[0], crate::Error::KeyError(error) + if error.abort == expected_message)); + } + error => panic!("unexpected error: {error:?}"), + } + let mut keys = lock_keys.lock().unwrap().clone(); + keys.sort(); + assert_eq!(keys, vec![vec![1], vec![11]]); + assert_eq!(*rollback_keys.lock().unwrap(), vec![vec![vec![1]]]); + } + #[rstest::rstest] #[case(Keyspace::Disable)] #[case(Keyspace::Enable { keyspace_id: 0 })] diff --git a/src/transaction/transaction/scanner.rs b/src/transaction/transaction/scanner.rs index 67959f30..a8e3296d 100644 --- a/src/transaction/transaction/scanner.rs +++ b/src/transaction/transaction/scanner.rs @@ -11,19 +11,21 @@ use tokio::time::sleep; use super::Transaction; use super::TransactionStatus; use crate::backoff::Backoff; +use crate::common::is_grpc_error; use crate::pd::PdClient; use crate::pd::PdRpcClient; use crate::proto::pdpb::Timestamp; use crate::region::RegionWithLeader; use crate::request::plan::handle_region_error; -use crate::request::plan::is_grpc_error; use crate::request::EncodeKeyspace; use crate::request::KeyMode; use crate::request::Keyspace; use crate::request::Plan; use crate::request::PlanBuilder; +use crate::request::PlanContext; use crate::request::RetryOptions; use crate::request::TruncateKeyspace; +use crate::store::reload_txn_protocol_range; use crate::store::HasKeyErrors; use crate::store::HasRegionError; use crate::transaction::buffer::MutationIterator; @@ -129,7 +131,7 @@ struct MutationCursor<'a> { struct RemoteScanner { status: Arc, pd_client: Arc, - keyspace: Keyspace, + plan_context: PlanContext, timestamp: Timestamp, retry_options: RetryOptions, end: Option, @@ -146,7 +148,7 @@ impl<'a, PdC: PdClient> Scanner<'a, PdC> { let range_empty = end .as_ref() .is_some_and(|end| !end.is_empty() && &start >= end); - let keyspace = txn.keyspace; + let keyspace = txn.plan_context.keyspace(); let remote = RemoteScanner::new(txn, start.clone(), end.clone(), range_empty); let mutations = if range_empty { None @@ -262,7 +264,7 @@ impl RemoteScanner { RemoteScanner { status: txn.status.clone(), pd_client: txn.rpc.clone(), - keyspace: txn.keyspace, + plan_context: txn.plan_context, timestamp: txn.timestamp.clone(), retry_options: txn.options.retry_options.clone(), end, @@ -320,7 +322,7 @@ impl RemoteScanner { /// page starts at the region's end key. async fn fetch_remote_page(&mut self) -> Result> { self.check_allow_operation()?; - let keyspace = self.keyspace; + let keyspace = self.plan_context.keyspace(); let mut region_backoff = self.retry_options.region_backoff.clone(); loop { @@ -378,15 +380,19 @@ impl RemoteScanner { false, false, ); - let plan = match PlanBuilder::new(self.pd_client.clone(), keyspace, request) - .single_region_with_store(region_store.clone()) - .await + let requirement = crate::store::Request::txn_protocol_requirement(&request); + let plan = match PlanBuilder::new_with_context( + self.pd_client.clone(), + request, + self.plan_context, + ) + .single_region_with_store(region_store.clone()) + .await { Ok(builder) => builder .resolve_lock( self.timestamp.clone(), self.retry_options.lock_backoff.clone(), - keyspace, ) .plan(), Err(error) => { @@ -407,8 +413,27 @@ impl RemoteScanner { }; if let Some(region_error) = response.region_error() { + if region_error.undetermined_result.is_none() { + if let Some(incompatible) = region_error.incompatible_request.as_ref() { + if reload_txn_protocol_range( + self.pd_client.clone(), + region_store.store_id, + region_store.txn_protocol_version_range, + self.plan_context.default_txn_protocol_version(), + requirement, + incompatible, + &mut region_backoff, + ) + .await + .is_some() + { + continue; + } + return Err(Error::IncompatibleRequest(Box::new(incompatible.clone()))); + } + } let Some(delay) = region_backoff.next_delay_duration() else { - return Err(Error::RegionError(Box::new(region_error))); + return Err(Error::from(region_error)); }; let retry_immediately = handle_region_error(self.pd_client.clone(), region_error, region_store).await?; @@ -454,6 +479,9 @@ impl RemoteScanner { backoff: &mut Backoff, error: Error, ) -> Result<()> { + if matches!(error, Error::IncompatibleRequest(_)) { + return Err(error); + } self.pd_client .invalidate_region_cache(region.ver_id()) .await; @@ -519,6 +547,7 @@ mod tests { use crate::request::Keyspace; use crate::store::RegionStore; use crate::store::Store; + use crate::store::TxnProtocolVersionRange; use crate::Key; use crate::KvPair; @@ -552,6 +581,105 @@ mod tests { )); } + #[rstest::rstest] + #[case(crate::mock::ReloadOutcome::Changed)] + #[case(crate::mock::ReloadOutcome::Unchanged)] + #[tokio::test] + async fn scanner_admission_reload_reselects_before_resending( + #[case] outcome: crate::mock::ReloadOutcome, + ) { + let versions = Arc::new(Mutex::new(Vec::new())); + let recorded = versions.clone(); + let incompatible = crate::mock::upper_admission_rejection(); + let rejection = incompatible.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + let request = request.downcast_ref::().unwrap(); + let version = request.context.as_ref().unwrap().txn_protocol_version; + recorded.lock().unwrap().push(version); + Ok(Box::new(kvrpcpb::ScanResponse { + error: (version == 1).then_some(kvrpcpb::KeyError { + locked: Some(kvrpcpb::LockInfo { + key: b"k".to_vec(), + primary_lock: b"k".to_vec(), + lock_version: 1, + ..Default::default() + }), + ..Default::default() + }), + region_error: (version == 1).then_some(errorpb::Error { + incompatible_request: Some(rejection.clone()), + ..Default::default() + }), + ..Default::default() + }) as Box) + }); + let fixture = crate::mock::ProtocolTestClient::new(client, outcome); + let mut txn = Transaction::new( + Timestamp::default(), + fixture.pd_client, + TransactionOptions::new_optimistic().read_only(), + Keyspace::Disable, + ); + let result = txn.scanner("k".to_owned()..="z".to_owned()).await; + if matches!(outcome, crate::mock::ReloadOutcome::Changed) { + assert!(result.is_ok()); + assert_eq!(*versions.lock().unwrap(), vec![1, 0]); + } else { + assert!( + matches!(result, Err(crate::Error::IncompatibleRequest(error)) if *error == incompatible) + ); + assert_eq!(*versions.lock().unwrap(), vec![1]); + } + assert_eq!(*fixture.reloads.lock().unwrap(), vec![42]); + } + + #[rstest::rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn scanner_region_error_bypasses_lock_resolution(#[case] undetermined: bool) { + let attempts = Arc::new(AtomicUsize::new(0)); + let recorded = attempts.clone(); + let original = errorpb::Error { + message: "original scan region error".into(), + undetermined_result: undetermined.then_some(Default::default()), + epoch_not_match: (!undetermined).then_some(Default::default()), + ..Default::default() + }; + let rejection = original.clone(); + let client = MockKvClient::with_dispatch_hook(move |request| { + assert!( + request.is::(), + "region errors must bypass auxiliary RPCs" + ); + recorded.fetch_add(1, Ordering::SeqCst); + Ok(Box::new(kvrpcpb::ScanResponse { + region_error: Some(rejection.clone()), + error: Some(kvrpcpb::KeyError { + locked: Some(kvrpcpb::LockInfo { + key: b"k".to_vec(), + primary_lock: b"k".to_vec(), + lock_version: 1, + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }) as Box) + }); + let mut txn = Transaction::new( + Timestamp::default(), + Arc::new(MockPdClient::new(client)), + TransactionOptions::new_optimistic() + .read_only() + .no_resolve_regions(), + Keyspace::Disable, + ); + assert!(matches!(txn.scanner("k".to_owned()..="z".to_owned()).await, + Err(crate::Error::RegionError(error)) if *error == original)); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + } + fn scanner_test_region( id: RegionId, start_key: Vec, @@ -658,6 +786,46 @@ mod tests { async fn invalidate_store_cache(&self, _store_id: u64) {} } + #[tokio::test] + async fn scanner_returns_local_incompatibility_without_retrying() { + let mappings = Arc::new(AtomicUsize::new(0)); + let invalidations = Arc::new(AtomicUsize::new(0)); + let kv_client = MockKvClient::with_dispatch_hook(|_req: &dyn Any| { + panic!("an incompatible request must not be dispatched"); + }); + let mapped_client = kv_client.clone(); + let mapped_count = mappings.clone(); + let invalidated_count = invalidations.clone(); + let pd_client = Arc::new( + MockPdClient::new(kv_client) + .with_map_region_to_store_hook(move |region| { + mapped_count.fetch_add(1, Ordering::SeqCst); + Ok(RegionStore::with_metadata( + region.clone(), + region.get_store_id().unwrap(), + TxnProtocolVersionRange { min: 2, max: 1 }, + Arc::new(mapped_client.clone()), + )) + }) + .with_invalidate_region_hook(move |_| { + invalidated_count.fetch_add(1, Ordering::SeqCst); + }), + ); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic().read_only(), + Keyspace::Disable, + ); + + assert!(matches!( + txn.scanner("k".to_owned()..="z".to_owned()).await, + Err(crate::Error::IncompatibleRequest(_)) + )); + assert_eq!(mappings.load(Ordering::SeqCst), 1); + assert_eq!(invalidations.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn scanner_paginates_lazily_after_eager_initial_page() { // One pair beyond the batch boundary is enough to exercise pagination.