diff --git a/src/pd/timestamp.rs b/src/pd/timestamp.rs index 65150d7e..d9daee10 100644 --- a/src/pd/timestamp.rs +++ b/src/pd/timestamp.rs @@ -3,8 +3,6 @@ //! Low-level mechanisms for obtaining timestamps from PD or the TSO //! microservice. -use std::collections::hash_map::Entry; -use std::collections::HashMap; use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; @@ -37,20 +35,57 @@ const MAX_BATCH_SIZE: usize = 64; /// TODO: This value should be adjustable. const MAX_PENDING_COUNT: usize = 1 << 16; +// Bound retained connections while allowing in-flight requests to finish after eviction. +const MAX_CACHED_API_V3_ORACLES: usize = 64; + struct TimestampRequest { sender: oneshot::Sender, } -#[derive(Clone)] struct ApiV3TimestampOracle { request_tx: mpsc::Sender, + task: tokio::task::JoinHandle>, +} + +impl Drop for ApiV3TimestampOracle { + fn drop(&mut self) { + self.task.abort(); + } +} + +#[derive(Default)] +struct ApiV3OracleCache { + entries: VecDeque<((u32, u32), Arc)>, +} + +impl ApiV3OracleCache { + fn get(&mut self, key: (u32, u32)) -> Option> { + let index = self + .entries + .iter() + .position(|(identity, _)| *identity == key)?; + let entry = self.entries.remove(index).unwrap(); + if entry.1.request_tx.is_closed() { + return None; + } + let oracle = entry.1.clone(); + self.entries.push_back(entry); + Some(oracle) + } + + fn insert(&mut self, key: (u32, u32), oracle: Arc) { + if self.entries.len() == MAX_CACHED_API_V3_ORACLES { + self.entries.pop_front(); + } + self.entries.push_back((key, oracle)); + } } struct ApiV3TimestampOracles { cluster_id: u64, pd_client: PdClient, security_mgr: Arc, - oracles: Mutex>, + oracles: Mutex, } /// The timestamp oracle (TSO) which provides monotonically increasing timestamps. @@ -80,7 +115,7 @@ impl TimestampOracle { cluster_id, pd_client: pd_client.clone(), security_mgr, - oracles: Mutex::new(HashMap::new()), + oracles: Mutex::new(ApiV3OracleCache::default()), }), }) } @@ -94,21 +129,23 @@ impl TimestampOracle { identity: Option, ) -> Result { debug!("getting current timestamp"); - let request_tx = match identity { - Some(identity) => { - self.api_v3_oracles - .get_or_create(identity) - .await? - .request_tx - } - None => self.legacy_request_tx, + let oracle = match identity { + Some(identity) => Some(self.api_v3_oracles.get_or_create(identity).await?), + None => None, }; + let request_tx = oracle + .as_ref() + .map(|oracle| oracle.request_tx.clone()) + .unwrap_or(self.legacy_request_tx); let (sender, response) = oneshot::channel(); request_tx .send(TimestampRequest { sender }) .await .map_err(|_| internal_err!("TimestampRequest channel is closed"))?; - Ok(response.await?) + let result = response.await; + // Keep the stream alive through the response even if its cache entry was evicted. + drop(oracle); + Ok(result?) } } @@ -116,18 +153,19 @@ impl ApiV3TimestampOracles { async fn get_or_create( &self, identity: apipb::KeyspaceIdentity, - ) -> Result { + ) -> Result> { let key = (identity.namespace_id, identity.keyspace_id); - if let Some(oracle) = self.oracles.lock().await.get(&key).cloned() { + if let Some(oracle) = self.oracles.lock().await.get(key) { return Ok(oracle); } - let oracle = self.connect(identity).await?; + let oracle = Arc::new(self.connect(identity).await?); let mut oracles = self.oracles.lock().await; - Ok(match oracles.entry(key) { - Entry::Occupied(entry) => entry.get().clone(), - Entry::Vacant(entry) => entry.insert(oracle).clone(), - }) + if let Some(existing) = oracles.get(key) { + return Ok(existing); + } + oracles.insert(key, oracle.clone()); + Ok(oracle) } async fn connect(&self, identity: apipb::KeyspaceIdentity) -> Result { @@ -167,7 +205,7 @@ impl ApiV3TimestampOracles { .await?; let callee_id = callee_id(&primary_url); let (request_tx, request_rx) = mpsc::channel(MAX_BATCH_SIZE); - tokio::spawn(run_api_v3_tso( + let task = tokio::spawn(run_api_v3_tso( self.cluster_id, keyspace_group_id, callee_id, @@ -175,7 +213,7 @@ impl ApiV3TimestampOracles { tso_client, request_rx, )); - return Ok(ApiV3TimestampOracle { request_tx }); + return Ok(ApiV3TimestampOracle { request_tx, task }); } Err(error) => last_error = Some(error), } @@ -281,7 +319,7 @@ async fn run_api_v3_tso( mut tso_client: tsopb::tso_client::TsoClient, request_rx: mpsc::Receiver, ) -> Result<()> { - let pending_requests = Arc::new(Mutex::new(VecDeque::with_capacity(MAX_PENDING_COUNT))); + let pending_requests = Arc::new(Mutex::new(VecDeque::new())); let sending_future_waker = Arc::new(AtomicWaker::new()); let request_stream = ApiV3TsoRequestStream { cluster_id, @@ -466,6 +504,71 @@ fn allocate_timestamps( mod tests { use super::*; + fn idle_oracle() -> Arc { + let (request_tx, mut request_rx) = mpsc::channel::(1); + let task = tokio::spawn(async move { + while let Some(request) = request_rx.recv().await { + let _ = request.sender.send(Timestamp::default()); + } + Ok(()) + }); + Arc::new(ApiV3TimestampOracle { request_tx, task }) + } + + #[tokio::test] + async fn api_v3_cache_evicts_least_recently_used_stream() { + let mut cache = ApiV3OracleCache::default(); + let first = idle_oracle(); + let second = idle_oracle(); + let first_task = first.task.abort_handle(); + let second_task = second.task.abort_handle(); + cache.insert((7, 1), first); + cache.insert((19, 1), second); + for id in 2..MAX_CACHED_API_V3_ORACLES as u32 { + cache.insert((7, id), idle_oracle()); + } + drop(cache.get((7, 1)).unwrap()); + cache.insert((7, 100), idle_oracle()); + tokio::task::yield_now().await; + assert_eq!(cache.entries.len(), MAX_CACHED_API_V3_ORACLES); + assert!(!first_task.is_finished()); + assert!(second_task.is_finished()); + assert!(cache.get((19, 1)).is_none()); + } + + #[tokio::test] + async fn api_v3_eviction_keeps_inflight_request_alive() { + let mut cache = ApiV3OracleCache::default(); + let (request_tx, mut request_rx) = mpsc::channel::(1); + let (release, ready) = oneshot::channel(); + let task = tokio::spawn(async move { + let request = request_rx.recv().await.unwrap(); + ready.await.unwrap(); + let _ = request.sender.send(Timestamp::default()); + futures::future::pending::<()>().await; + Ok(()) + }); + let task_handle = task.abort_handle(); + let active = Arc::new(ApiV3TimestampOracle { request_tx, task }); + cache.insert((7, 1), active.clone()); + let (sender, response) = oneshot::channel(); + active + .request_tx + .send(TimestampRequest { sender }) + .await + .unwrap(); + for id in 2..=MAX_CACHED_API_V3_ORACLES as u32 + 1 { + cache.insert((7, id), idle_oracle()); + } + assert!(cache.get((7, 1)).is_none()); + assert!(!task_handle.is_finished()); + release.send(()).unwrap(); + response.await.unwrap(); + drop(active); + tokio::task::yield_now().await; + assert!(task_handle.is_finished()); + } + #[tokio::test] async fn api_v3_tso_request_stream_batches_and_includes_identity() { let (request_tx, request_rx) = mpsc::channel(2); diff --git a/src/request/keyspace.rs b/src/request/keyspace.rs index 305f0d98..f75f9804 100644 --- a/src/request/keyspace.rs +++ b/src/request/keyspace.rs @@ -11,7 +11,7 @@ use crate::{BoundRange, KvPair}; pub const RAW_KEY_PREFIX: u8 = b'r'; pub const TXN_KEY_PREFIX: u8 = b'x'; pub const KEYSPACE_PREFIX_LEN: usize = 4; -pub const API_V3_PREFIX_LEN: usize = KEYSPACE_PREFIX_LEN; +pub const API_V3_PREFIX_LEN: usize = 8; pub const API_V3_MAX_KEYSPACE_ID: u32 = 0xFF_FFFF; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -104,13 +104,13 @@ impl Keyspace { pub fn v3_route_prefix(&self, key_mode: KeyMode) -> Option<[u8; API_V3_PREFIX_LEN]> { let Keyspace::ApiV3 { - namespace_id: _, + namespace_id, keyspace_id, } = *self else { return None; }; - Some(api_v3_keyspace_prefix(keyspace_id, key_mode)) + Some(api_v3_keyspace_prefix(namespace_id, keyspace_id, key_mode)) } pub fn encode_route_key(&self, key: &Key, key_mode: KeyMode) -> Key { @@ -184,13 +184,13 @@ impl Keyspace { fn v3_route_range_end(&self, key_mode: KeyMode) -> Option> { let Keyspace::ApiV3 { - namespace_id: _, + namespace_id, keyspace_id, } = *self else { return None; }; - let start = u32::from_be_bytes(api_v3_keyspace_prefix(keyspace_id, key_mode)); + let start = u64::from_be_bytes(api_v3_keyspace_prefix(namespace_id, keyspace_id, key_mode)); Some(start.wrapping_add(1).to_be_bytes().to_vec()) } } @@ -209,6 +209,8 @@ pub(crate) fn keyspace_meta_identity( pub(crate) fn keyspace_meta_legacy_id(keyspace: &keyspacepb::KeyspaceMeta) -> Option { match keyspace.keyspace { Some(keyspacepb::keyspace_meta::Keyspace::Id(id)) => Some(id), + // Legacy PD uses a proto3 scalar, so keyspace ID zero is omitted on the wire. + None => Some(0), _ => None, } } @@ -441,18 +443,20 @@ fn keyspace_prefix(keyspace_id: u32, key_mode: KeyMode) -> [u8; KEYSPACE_PREFIX_ prefix } -fn api_v3_keyspace_prefix(keyspace_id: u32, key_mode: KeyMode) -> [u8; API_V3_PREFIX_LEN] { +fn api_v3_keyspace_prefix( + namespace_id: u32, + keyspace_id: u32, + key_mode: KeyMode, +) -> [u8; API_V3_PREFIX_LEN] { debug_assert!(keyspace_id <= API_V3_MAX_KEYSPACE_ID); - let keyspace_bytes = keyspace_id.to_be_bytes(); - [ - match key_mode { - KeyMode::Raw => RAW_KEY_PREFIX, - KeyMode::Txn => TXN_KEY_PREFIX, - }, - keyspace_bytes[1], - keyspace_bytes[2], - keyspace_bytes[3], - ] + let mut prefix = [0; API_V3_PREFIX_LEN]; + prefix[0] = match key_mode { + KeyMode::Raw => RAW_KEY_PREFIX, + KeyMode::Txn => TXN_KEY_PREFIX, + }; + prefix[1..5].copy_from_slice(&namespace_id.to_be_bytes()); + prefix[5..].copy_from_slice(&keyspace_id.to_be_bytes()[1..]); + prefix } fn prepend_bytes(vec: &mut Vec, prefix: &[u8; N]) { @@ -476,6 +480,35 @@ fn pretruncate_bytes(vec: &mut Vec) { mod tests { use super::*; + #[test] + fn test_legacy_keyspace_metadata_wire_compatibility() { + use prost::Message; + + // Legacy proto3 metadata for DEFAULT omits the scalar ID zero. + let default = keyspacepb::KeyspaceMeta::decode(&b"\x12\x07DEFAULT"[..]).unwrap(); + assert_eq!(keyspace_meta_legacy_id(&default), Some(0)); + assert_eq!(keyspace_meta_identity(&default), None); + + let nonzero = keyspacepb::KeyspaceMeta::decode(&b"\x08\x2a\x12\x04test"[..]).unwrap(); + assert_eq!(keyspace_meta_legacy_id(&nonzero), Some(42)); + + let scoped = keyspacepb::KeyspaceMeta { + keyspace: Some(keyspacepb::keyspace_meta::Keyspace::KeyspaceIdentity( + apipb::KeyspaceIdentity { + namespace_id: 7, + keyspace_id: 0, + }, + )), + ..Default::default() + }; + let decoded = keyspacepb::KeyspaceMeta::decode(scoped.encode_to_vec().as_slice()).unwrap(); + assert_eq!(keyspace_meta_legacy_id(&decoded), None); + assert_eq!( + keyspace_meta_identity(&decoded), + keyspace_meta_identity(&scoped) + ); + } + #[test] fn test_keyspace_prefix() { let key_mode = KeyMode::Raw; @@ -666,25 +699,31 @@ mod tests { assert_eq!( keyspace.encode_route_key(&Key::from(vec![b'k']), KeyMode::Raw), - Key::from(vec![b'r', 5, 6, 7, b'k']) + Key::from(vec![b'r', 1, 2, 3, 4, 5, 6, 7, b'k']) ); assert_eq!( keyspace.encode_route_key(&Key::from(vec![b'k']), KeyMode::Txn), - Key::from(vec![b'x', 5, 6, 7, b'k']) + Key::from(vec![b'x', 1, 2, 3, 4, 5, 6, 7, b'k']) ); let route_range = keyspace.encode_route_range(Key::from(vec![b'a']), Key::from(vec![b'z']), KeyMode::Txn); - assert_eq!(route_range.0, Key::from(vec![b'x', 5, 6, 7, b'a'])); - assert_eq!(route_range.1, Key::from(vec![b'x', 5, 6, 7, b'z'])); + assert_eq!( + route_range.0, + Key::from(vec![b'x', 1, 2, 3, 4, 5, 6, 7, b'a']) + ); + assert_eq!( + route_range.1, + Key::from(vec![b'x', 1, 2, 3, 4, 5, 6, 7, b'z']) + ); assert_eq!( keyspace.decode_route_range(route_range.0, route_range.1, KeyMode::Txn), (Key::from(vec![b'a']), Key::from(vec![b'z'])) ); let whole_range = keyspace.encode_route_range(Key::EMPTY, Key::EMPTY, KeyMode::Txn); - assert_eq!(whole_range.0, Key::from(vec![b'x', 5, 6, 7])); - assert_eq!(whole_range.1, Key::from(vec![b'x', 5, 6, 8])); + assert_eq!(whole_range.0, Key::from(vec![b'x', 1, 2, 3, 4, 5, 6, 7])); + assert_eq!(whole_range.1, Key::from(vec![b'x', 1, 2, 3, 4, 5, 6, 8])); assert_eq!( keyspace.decode_route_range(whole_range.0, whole_range.1, KeyMode::Txn), (Key::EMPTY, Key::EMPTY) @@ -693,14 +732,38 @@ mod tests { let last_keyspace = Keyspace::api_v3(u32::MAX, API_V3_MAX_KEYSPACE_ID).unwrap(); let last_whole_range = last_keyspace.encode_route_range(Key::EMPTY, Key::EMPTY, KeyMode::Txn); - assert_eq!(last_whole_range.0, Key::from(vec![b'x', 0xff, 0xff, 0xff])); - assert_eq!(last_whole_range.1, Key::from(vec![b'y', 0, 0, 0])); + assert_eq!( + last_whole_range.0, + Key::from(vec![b'x', 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff]) + ); + assert_eq!( + last_whole_range.1, + Key::from(vec![b'y', 0, 0, 0, 0, 0, 0, 0]) + ); assert_eq!( last_keyspace.decode_route_range(last_whole_range.0, last_whole_range.1, KeyMode::Txn), (Key::EMPTY, Key::EMPTY) ); } + #[test] + fn test_api_v3_namespace_isolation_and_range_carry() { + let first = Keyspace::api_v3(1, 7).unwrap(); + let second = Keyspace::api_v3(2, 7).unwrap(); + let logical = Key::from(vec![b'x', 0, 0, 0, 1, 0, 0, 7]); + let first_key = first.encode_route_key(&logical, KeyMode::Txn); + let second_key = second.encode_route_key(&logical, KeyMode::Txn); + assert_ne!(first_key, second_key); + assert_eq!(first_key.0.len(), 16); + assert_eq!( + first.decode_route_range(first_key.clone(), first_key, KeyMode::Txn), + (logical.clone(), logical) + ); + let last = Keyspace::api_v3(1, API_V3_MAX_KEYSPACE_ID).unwrap(); + let range = last.encode_route_range(Key::EMPTY, Key::EMPTY, KeyMode::Txn); + assert_eq!(range.1, Key::from(vec![b'x', 0, 0, 0, 2, 0, 0, 0])); + } + #[test] fn test_api_v3_wire_key_is_noop() { let keyspace = Keyspace::api_v3(7, 0xDEAD).unwrap(); diff --git a/src/transaction/client.rs b/src/transaction/client.rs index 6f744251..80266183 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -215,6 +215,14 @@ impl Client { Ok(Client { pd, keyspace }) } + /// Returns a client for an exact API V3 identity, sharing the PD connection and region cache. + pub fn with_keyspace_identity(&self, namespace_id: u32, keyspace_id: u32) -> Result { + Ok(Client { + pd: self.pd.clone(), + keyspace: Keyspace::api_v3(namespace_id, keyspace_id)?, + }) + } + /// Create a transactional [`Client`] that uses API V2 without adding or removing any API V2 /// keyspace/key-mode prefix, with a custom configuration. ///