diff --git a/src/pd/client.rs b/src/pd/client.rs index 14fea6b9..88d2976f 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -76,6 +76,20 @@ pub trait PdClient: Send + Sync + 'static { async fn update_safepoint(self: Arc, safepoint: u64) -> Result; + /// Updates only the supplied V3 identity, without falling back to global GC. + async fn update_safepoint_with_identity( + self: Arc, + safepoint: u64, + identity: Option, + ) -> Result { + if identity.is_some() { + return Err(crate::internal_err!( + "scoped GC safe point updates are not supported by this PD client" + )); + } + self.update_safepoint(safepoint).await + } + async fn load_keyspace(&self, keyspace: &str) -> Result; async fn lookup_keyspaces(&self, keyspace: &str) -> Result>; @@ -289,6 +303,17 @@ impl PdClient for PdRpcClient { self.pd.clone().update_safepoint(safepoint).await } + async fn update_safepoint_with_identity( + self: Arc, + safepoint: u64, + identity: Option, + ) -> Result { + self.pd + .clone() + .update_safepoint_with_identity(safepoint, identity) + .await + } + async fn update_leader(&self, ver_id: RegionVerId, leader: metapb::Peer) -> Result<()> { self.region_cache.update_leader(ver_id, leader).await } diff --git a/src/pd/cluster.rs b/src/pd/cluster.rs index 214b10c2..bccb5f8b 100644 --- a/src/pd/cluster.rs +++ b/src/pd/cluster.rs @@ -103,6 +103,26 @@ impl Cluster { req.send(&mut self.client, timeout).await } + pub async fn update_safepoint_with_identity( + &mut self, + safepoint: u64, + identity: Option, + timeout: Duration, + ) -> Result { + let Some(identity) = identity else { + return self + .update_safepoint(safepoint, timeout) + .await + .map(|resp| resp.new_safe_point == safepoint); + }; + let mut req = pd_request!(self.id, pdpb::UpdateGcSafePointV2Request); + req.safe_point = safepoint; + req.keyspace = + Some(pdpb::update_gc_safe_point_v2_request::Keyspace::KeyspaceIdentity(identity)); + let resp = req.send(&mut self.client, timeout).await?; + Ok(resp.new_safe_point == safepoint) + } + pub async fn load_keyspace( &mut self, keyspace: &str, @@ -463,6 +483,16 @@ impl PdMessage for pdpb::UpdateGcSafePointRequest { } } +#[async_trait] +impl PdMessage for pdpb::UpdateGcSafePointV2Request { + type Client = pdpb::pd_client::PdClient; + type Response = pdpb::UpdateGcSafePointV2Response; + + async fn rpc(req: Request, client: &mut Self::Client) -> GrpcResult { + Ok(client.update_gc_safe_point_v2(req).await?.into_inner()) + } +} + #[async_trait] impl PdMessage for keyspacepb::LoadKeyspaceRequest { type Client = keyspacepb::keyspace_client::KeyspaceClient; @@ -511,6 +541,12 @@ impl PdResponse for pdpb::UpdateGcSafePointResponse { } } +impl PdResponse for pdpb::UpdateGcSafePointV2Response { + fn header(&self) -> &pdpb::ResponseHeader { + self.header.as_ref().unwrap() + } +} + impl PdResponse for keyspacepb::LoadKeyspaceResponse { fn header(&self) -> &pdpb::ResponseHeader { self.header.as_ref().unwrap() @@ -522,3 +558,6 @@ impl PdResponse for keyspacepb::LookupKeyspaceResponse { self.header.as_ref().unwrap() } } + +#[cfg(test)] +mod gc_tests; diff --git a/src/pd/cluster/gc_tests.rs b/src/pd/cluster/gc_tests.rs new file mode 100644 index 00000000..9ab1d42c --- /dev/null +++ b/src/pd/cluster/gc_tests.rs @@ -0,0 +1,241 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +use std::convert::Infallible; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::Mutex; +use std::task::{Context, Poll}; + +use tonic::body::BoxBody; +use tonic::codegen::{http, BoxFuture, Service}; +use tonic::server::NamedService; + +use super::*; + +type GcCall = (Option, u64); + +#[derive(Clone, Default)] +struct GcService { + calls: Arc>>, + // 0: accepted; 1: already newer; 2: PD header error; 3: unsupported RPC. + reply: Arc, +} + +impl GcService { + fn response(&self, safe_point: u64) -> GrpcResult<(pdpb::ResponseHeader, u64)> { + let mut header = pdpb::ResponseHeader { + cluster_id: 42, + ..Default::default() + }; + match self.reply.load(Ordering::SeqCst) { + 1 => Ok((header, safe_point + 1)), + 2 => { + header.error = Some(pdpb::Error { + message: "scoped GC rejected".into(), + ..Default::default() + }); + Ok((header, safe_point)) + } + 3 => Err(tonic::Status::unimplemented("scoped GC unavailable")), + _ => Ok((header, safe_point)), + } + } +} + +struct ScopedGc(GcService); + +impl tonic::server::UnaryService for ScopedGc { + type Response = pdpb::UpdateGcSafePointV2Response; + type Future = BoxFuture, tonic::Status>; + + fn call(&mut self, request: Request) -> Self::Future { + let service = self.0.clone(); + Box::pin(async move { + let request = request.into_inner(); + assert_eq!(request.header.unwrap().cluster_id, 42); + let Some(pdpb::update_gc_safe_point_v2_request::Keyspace::KeyspaceIdentity(identity)) = + request.keyspace + else { + panic!("V3 GC must send the complete identity"); + }; + service + .calls + .lock() + .unwrap() + .push((Some(identity), request.safe_point)); + let (header, new_safe_point) = service.response(request.safe_point)?; + Ok(tonic::Response::new(pdpb::UpdateGcSafePointV2Response { + header: Some(header), + new_safe_point, + })) + }) + } +} + +struct GlobalGc(GcService); + +impl tonic::server::UnaryService for GlobalGc { + type Response = pdpb::UpdateGcSafePointResponse; + type Future = BoxFuture, tonic::Status>; + + fn call(&mut self, request: Request) -> Self::Future { + let service = self.0.clone(); + Box::pin(async move { + let request = request.into_inner(); + assert_eq!(request.header.unwrap().cluster_id, 42); + service + .calls + .lock() + .unwrap() + .push((None, request.safe_point)); + let (header, new_safe_point) = service.response(request.safe_point)?; + Ok(tonic::Response::new(pdpb::UpdateGcSafePointResponse { + header: Some(header), + new_safe_point, + })) + }) + } +} + +impl Service> for GcService { + type Response = http::Response; + type Error = Infallible; + type Future = BoxFuture; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + let service = self.clone(); + Box::pin(async move { + match request.uri().path() { + "/pdpb.PD/UpdateGCSafePointV2" => { + let codec = tonic::codec::ProstCodec::default(); + Ok(tonic::server::Grpc::new(codec) + .unary(ScopedGc(service), request) + .await) + } + "/pdpb.PD/UpdateGCSafePoint" => { + let codec = tonic::codec::ProstCodec::default(); + Ok(tonic::server::Grpc::new(codec) + .unary(GlobalGc(service), request) + .await) + } + _ => Ok(http::Response::builder() + .header("grpc-status", "12") + .header("content-type", "application/grpc") + .body(tonic::body::empty_body()) + .unwrap()), + } + }) + } +} + +impl NamedService for GcService { + const NAME: &'static str = "pdpb.PD"; +} + +#[tokio::test] +async fn gc_safepoint_rpc_preserves_scope_and_never_falls_back() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let service = GcService::default(); + let incoming = futures::stream::unfold(listener, |listener| async move { + let result = listener.accept().await.map(|(stream, _)| stream); + Some((result, listener)) + }); + let server = tokio::spawn( + tonic::transport::Server::builder() + .add_service(service.clone()) + .serve_with_incoming(incoming), + ); + let channel = Channel::from_shared(format!("http://{address}")) + .unwrap() + .connect() + .await + .unwrap(); + let client = pdpb::pd_client::PdClient::new(channel.clone()); + let tso = TimestampOracle::new(42, &client, Arc::new(SecurityManager::default())).unwrap(); + let mut cluster = Cluster { + id: 42, + client, + keyspace_client: keyspacepb::keyspace_client::KeyspaceClient::new(channel), + members: pdpb::GetMembersResponse::default(), + tso, + }; + let timeout = Duration::from_secs(3); + for namespace_id in [7, 19] { + let identity = apipb::KeyspaceIdentity { + namespace_id, + keyspace_id: 42, + }; + assert!(cluster + .update_safepoint_with_identity(100, Some(identity.clone()), timeout) + .await + .unwrap()); + assert_eq!( + service.calls.lock().unwrap().last(), + Some(&(Some(identity), 100)) + ); + } + for keyspace in [ + crate::request::Keyspace::Disable, + crate::request::Keyspace::Enable { keyspace_id: 42 }, + crate::request::Keyspace::ApiV2NoPrefix, + ] { + assert!(cluster + .update_safepoint_with_identity(100, keyspace.v3_identity(), timeout) + .await + .unwrap()); + } + let identity = apipb::KeyspaceIdentity { + namespace_id: 7, + keyspace_id: 42, + }; + service.reply.store(1, Ordering::SeqCst); + assert!(!cluster + .update_safepoint_with_identity(100, Some(identity.clone()), timeout) + .await + .unwrap()); + for reply in [2, 3] { + service.reply.store(reply, Ordering::SeqCst); + cluster + .update_safepoint_with_identity(100, Some(identity.clone()), timeout) + .await + .unwrap_err(); + } + let calls = service.calls.lock().unwrap(); + assert_eq!(calls.len(), 8); + assert_eq!( + calls + .iter() + .filter(|(identity, _)| identity.is_none()) + .count(), + 3 + ); + drop(calls); + server.abort(); +} + +#[tokio::test] +async fn unsupported_pd_implementation_does_not_fall_back_to_global_gc() { + use crate::pd::PdClient; + let pd = Arc::new(crate::mock::MockPdClient::new( + crate::mock::MockKvClient::default(), + )); + // The mock's global update panics; the default scoped method must reject + // the request without reaching it. + let error = pd + .update_safepoint_with_identity( + 100, + Some(apipb::KeyspaceIdentity { + namespace_id: 7, + keyspace_id: 42, + }), + ) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("scoped GC safe point updates are not supported")); +} diff --git a/src/pd/retry.rs b/src/pd/retry.rs index ffac44b4..f2cb3f59 100644 --- a/src/pd/retry.rs +++ b/src/pd/retry.rs @@ -53,6 +53,12 @@ pub trait RetryClientTrait { async fn update_safepoint(self: Arc, safepoint: u64) -> Result; + async fn update_safepoint_with_identity( + self: Arc, + safepoint: u64, + identity: Option, + ) -> Result; + async fn load_keyspace(&self, keyspace: &str) -> Result; async fn lookup_keyspaces(&self, keyspace: &str) -> Result>; @@ -225,6 +231,18 @@ impl RetryClientTrait for RetryClient { }) } + async fn update_safepoint_with_identity( + self: Arc, + safepoint: u64, + identity: Option, + ) -> Result { + retry_mut!(self, "update_gc_safepoint_with_identity", |cluster| async { + cluster + .update_safepoint_with_identity(safepoint, identity.clone(), self.timeout) + .await + }) + } + async fn load_keyspace(&self, keyspace: &str) -> Result { retry_mut!(self, "load_keyspace", |cluster| async { cluster.load_keyspace(keyspace, self.timeout).await diff --git a/src/region_cache.rs b/src/region_cache.rs index 81119262..39a924c8 100644 --- a/src/region_cache.rs +++ b/src/region_cache.rs @@ -367,6 +367,14 @@ mod test { todo!() } + async fn update_safepoint_with_identity( + self: Arc, + _safepoint: u64, + _identity: Option, + ) -> Result { + unimplemented!() + } + async fn load_keyspace(&self, _keyspace: &str) -> Result { unimplemented!() } diff --git a/src/transaction/client.rs b/src/transaction/client.rs index 80266183..7a2bebd8 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -365,6 +365,9 @@ impl Client { /// 1. resolving all locks with timestamp <= `safepoint` /// 2. updating PD's known safepoint /// + /// API V3 clients update only their namespace/keyspace safe point. Other + /// client modes retain the legacy global safe point behavior. + /// /// This is a simplified version of [GC in TiDB](https://docs.pingcap.com/tidb/stable/garbage-collection-overview). /// We skip the second step "delete ranges" which is an optimization for TiDB. pub async fn gc(&self, safepoint: Timestamp) -> Result { @@ -380,7 +383,7 @@ impl Client { let res: bool = self .pd .clone() - .update_safepoint(safepoint.version()) + .update_safepoint_with_identity(safepoint.version(), self.keyspace.v3_identity()) .await?; if !res { info!(