diff --git a/Cargo.lock b/Cargo.lock index d7ee941b3b..9adb6bb506 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1812,6 +1812,7 @@ dependencies = [ "common-telemetry", "common-time", "datatypes", + "etcd-client", "futures", "lazy_static", "prost", diff --git a/Cargo.toml b/Cargo.toml index 02e657ed0e..93c190dad5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,7 @@ datafusion-optimizer = { git = "https://github.com/waynexia/arrow-datafusion.git datafusion-physical-expr = { git = "https://github.com/waynexia/arrow-datafusion.git", rev = "63e52dde9e44cac4b1f6c6e6b6bf6368ba3bd323" } datafusion-sql = { git = "https://github.com/waynexia/arrow-datafusion.git", rev = "63e52dde9e44cac4b1f6c6e6b6bf6368ba3bd323" } datafusion-substrait = { git = "https://github.com/waynexia/arrow-datafusion.git", rev = "63e52dde9e44cac4b1f6c6e6b6bf6368ba3bd323" } +etcd-client = "0.11" futures = "0.3" futures-util = "0.3" greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "917ead6274b4dccbaf33c59a7360646ba2f285a9" } diff --git a/src/catalog/src/remote.rs b/src/catalog/src/remote.rs index a69c8b28a8..f60e574a1a 100644 --- a/src/catalog/src/remote.rs +++ b/src/catalog/src/remote.rs @@ -30,76 +30,3 @@ pub trait KvCacheInvalidator: Send + Sync { } pub type KvCacheInvalidatorRef = Arc; - -#[cfg(test)] -mod tests { - use std::any::Any; - - use async_stream::stream; - use common_meta::kv_backend::{Kv, KvBackend, ValueIter}; - - use crate::error::Error; - - struct MockKvBackend {} - - #[async_trait::async_trait] - impl KvBackend for MockKvBackend { - type Error = Error; - - fn range<'a, 'b>(&'a self, _key: &[u8]) -> ValueIter<'b, Error> - where - 'a: 'b, - { - Box::pin(stream!({ - for i in 0..3 { - yield Ok(Kv( - i.to_string().as_bytes().to_vec(), - i.to_string().as_bytes().to_vec(), - )) - } - })) - } - - async fn set(&self, _key: &[u8], _val: &[u8]) -> Result<(), Error> { - unimplemented!() - } - - async fn compare_and_set( - &self, - _key: &[u8], - _expect: &[u8], - _val: &[u8], - ) -> Result>>, Error> { - unimplemented!() - } - - async fn delete_range(&self, _key: &[u8], _end: &[u8]) -> Result<(), Error> { - unimplemented!() - } - - async fn move_value(&self, _from_key: &[u8], _to_key: &[u8]) -> Result<(), Error> { - unimplemented!() - } - - fn as_any(&self) -> &dyn Any { - self - } - } - - #[tokio::test] - async fn test_get() { - let backend = MockKvBackend {}; - - let result = backend.get(0.to_string().as_bytes()).await; - assert_eq!(0.to_string().as_bytes(), result.unwrap().unwrap().0); - - let result = backend.get(1.to_string().as_bytes()).await; - assert_eq!(1.to_string().as_bytes(), result.unwrap().unwrap().0); - - let result = backend.get(2.to_string().as_bytes()).await; - assert_eq!(2.to_string().as_bytes(), result.unwrap().unwrap().0); - - let result = backend.get(3.to_string().as_bytes()).await; - assert!(result.unwrap().is_none()); - } -} diff --git a/src/catalog/src/remote/client.rs b/src/catalog/src/remote/client.rs index 02d16926fc..cd15b7265d 100644 --- a/src/catalog/src/remote/client.rs +++ b/src/catalog/src/remote/client.rs @@ -17,15 +17,18 @@ use std::fmt::Debug; use std::sync::Arc; use std::time::Duration; -use async_stream::stream; use common_error::prelude::BoxedError; use common_meta::error::Error::{CacheNotGet, GetKvCache}; use common_meta::error::{CacheNotGetSnafu, Error, MetaSrvSnafu, Result}; -use common_meta::kv_backend::{Kv, KvBackend, KvBackendRef, ValueIter}; +use common_meta::kv_backend::{KvBackend, KvBackendRef, TxnService}; use common_meta::rpc::store::{ - CompareAndPutRequest, DeleteRangeRequest, MoveValueRequest, PutRequest, RangeRequest, + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, + BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, + DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, + RangeRequest, RangeResponse, }; -use common_telemetry::{info, timer}; +use common_meta::rpc::KeyValue; +use common_telemetry::timer; use meta_client::client::MetaClient; use moka::future::{Cache, CacheBuilder}; use snafu::{OptionExt, ResultExt}; @@ -37,24 +40,29 @@ const CACHE_MAX_CAPACITY: u64 = 10000; const CACHE_TTL_SECOND: u64 = 10 * 60; const CACHE_TTI_SECOND: u64 = 5 * 60; -pub type CacheBackendRef = Arc, Kv>>; +pub type CacheBackendRef = Arc, KeyValue>>; + pub struct CachedMetaKvBackend { kv_backend: KvBackendRef, cache: CacheBackendRef, + name: String, +} + +impl TxnService for CachedMetaKvBackend { + type Error = Error; } #[async_trait::async_trait] impl KvBackend for CachedMetaKvBackend { - type Error = Error; - - fn range<'a, 'b>(&'a self, key: &[u8]) -> ValueIter<'b, Error> - where - 'a: 'b, - { - self.kv_backend.range(key) + fn name(&self) -> &str { + &self.name } - async fn get(&self, key: &[u8]) -> Result> { + async fn range(&self, req: RangeRequest) -> Result { + self.kv_backend.range(req).await + } + + async fn get(&self, key: &[u8]) -> Result> { let _timer = timer!(METRIC_CATALOG_KV_GET); let init = async { @@ -81,8 +89,10 @@ impl KvBackend for CachedMetaKvBackend { }) } - async fn set(&self, key: &[u8], val: &[u8]) -> Result<()> { - let ret = self.kv_backend.set(key, val).await; + async fn put(&self, req: PutRequest) -> Result { + let key = &req.key.clone(); + + let ret = self.kv_backend.put(req).await; if ret.is_ok() { self.invalidate_key(key).await; @@ -91,8 +101,72 @@ impl KvBackend for CachedMetaKvBackend { ret } - async fn delete(&self, key: &[u8]) -> Result<()> { - let ret = self.kv_backend.delete_range(key, &[]).await; + async fn batch_put(&self, req: BatchPutRequest) -> Result { + let keys = req + .kvs + .iter() + .map(|kv| kv.key().to_vec()) + .collect::>(); + + let resp = self.kv_backend.batch_put(req).await; + + if resp.is_ok() { + for key in keys { + self.invalidate_key(&key).await; + } + } + + resp + } + + async fn delete_range(&self, mut req: DeleteRangeRequest) -> Result { + let prev_kv = req.prev_kv; + + req.prev_kv = true; + let resp = self.kv_backend.delete_range(req).await; + match resp { + Ok(mut resp) => { + for prev_kv in resp.prev_kvs.iter() { + self.invalidate_key(prev_kv.key()).await; + } + + if !prev_kv { + resp.prev_kvs = vec![]; + } + Ok(resp) + } + Err(e) => Err(e), + } + } + + async fn batch_delete(&self, mut req: BatchDeleteRequest) -> Result { + let prev_kv = req.prev_kv; + + req.prev_kv = true; + let resp = self.kv_backend.batch_delete(req).await; + match resp { + Ok(mut resp) => { + for prev_kv in resp.prev_kvs.iter() { + self.invalidate_key(prev_kv.key()).await; + } + + if !prev_kv { + resp.prev_kvs = vec![]; + } + Ok(resp) + } + Err(e) => Err(e), + } + } + + async fn batch_get(&self, req: BatchGetRequest) -> Result { + self.kv_backend.batch_get(req).await + } + + async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result { + let key = &req.key.clone(); + + let ret = self.kv_backend.compare_and_put(req).await; if ret.is_ok() { self.invalidate_key(key).await; @@ -101,28 +175,11 @@ impl KvBackend for CachedMetaKvBackend { ret } - async fn delete_range(&self, _key: &[u8], _end: &[u8]) -> Result<()> { - // TODO(fys): implement it - unimplemented!() - } + async fn move_value(&self, req: MoveValueRequest) -> Result { + let from_key = &req.from_key.clone(); + let to_key = &req.to_key.clone(); - async fn compare_and_set( - &self, - key: &[u8], - expect: &[u8], - val: &[u8], - ) -> Result>>> { - let ret = self.kv_backend.compare_and_set(key, expect, val).await; - - if ret.is_ok() { - self.invalidate_key(key).await; - } - - ret - } - - async fn move_value(&self, from_key: &[u8], to_key: &[u8]) -> Result<()> { - let ret = self.kv_backend.move_value(from_key, to_key).await; + let ret = self.kv_backend.move_value(req).await; if ret.is_ok() { self.invalidate_key(from_key).await; @@ -146,15 +203,8 @@ impl KvCacheInvalidator for CachedMetaKvBackend { impl CachedMetaKvBackend { pub fn new(client: Arc) -> Self { - let cache = Arc::new( - CacheBuilder::new(CACHE_MAX_CAPACITY) - .time_to_live(Duration::from_secs(CACHE_TTL_SECOND)) - .time_to_idle(Duration::from_secs(CACHE_TTI_SECOND)) - .build(), - ); let kv_backend = Arc::new(MetaKvBackend { client }); - - Self { kv_backend, cache } + Self::wrap(kv_backend) } pub fn wrap(kv_backend: KvBackendRef) -> Self { @@ -165,7 +215,12 @@ impl CachedMetaKvBackend { .build(), ); - Self { kv_backend, cache } + let name = format!("CachedKvBackend({})", kv_backend.name()); + Self { + kv_backend, + cache, + name, + } } pub fn cache(&self) -> &CacheBackendRef { @@ -178,108 +233,97 @@ pub struct MetaKvBackend { pub client: Arc, } +impl TxnService for MetaKvBackend { + type Error = Error; +} + /// Implement `KvBackend` trait for `MetaKvBackend` instead of opendal's `Accessor` since /// `MetaClient`'s range method can return both keys and values, which can reduce IO overhead /// comparing to `Accessor`'s list and get method. #[async_trait::async_trait] impl KvBackend for MetaKvBackend { - type Error = Error; - - fn range<'a, 'b>(&'a self, key: &[u8]) -> ValueIter<'b, Error> - where - 'a: 'b, - { - let key = key.to_vec(); - Box::pin(stream!({ - let mut resp = self - .client - .range(RangeRequest::new().with_prefix(key)) - .await - .map_err(BoxedError::new) - .context(MetaSrvSnafu)?; - let kvs = resp.take_kvs(); - for mut kv in kvs.into_iter() { - yield Ok(Kv(kv.take_key(), kv.take_value())) - } - })) + fn name(&self) -> &str { + "MetaKvBackend" } - async fn get(&self, key: &[u8]) -> Result> { + async fn range(&self, req: RangeRequest) -> Result { + self.client + .range(req) + .await + .map_err(BoxedError::new) + .context(MetaSrvSnafu) + } + + async fn get(&self, key: &[u8]) -> Result> { let mut response = self .client .range(RangeRequest::new().with_key(key)) .await .map_err(BoxedError::new) .context(MetaSrvSnafu)?; - Ok(response - .take_kvs() - .get_mut(0) - .map(|kv| Kv(kv.take_key(), kv.take_value()))) + Ok(response.take_kvs().get_mut(0).map(|kv| KeyValue { + key: kv.take_key(), + value: kv.take_value(), + })) } - async fn set(&self, key: &[u8], val: &[u8]) -> Result<()> { - let req = PutRequest::new() - .with_key(key.to_vec()) - .with_value(val.to_vec()); - let _ = self - .client + async fn batch_put(&self, req: BatchPutRequest) -> Result { + self.client + .batch_put(req) + .await + .map_err(BoxedError::new) + .context(MetaSrvSnafu) + } + + async fn put(&self, req: PutRequest) -> Result { + self.client .put(req) .await .map_err(BoxedError::new) - .context(MetaSrvSnafu)?; - Ok(()) + .context(MetaSrvSnafu) } - async fn delete_range(&self, key: &[u8], end: &[u8]) -> Result<()> { - let req = DeleteRangeRequest::new().with_range(key.to_vec(), end.to_vec()); - let resp = self - .client + async fn delete_range(&self, req: DeleteRangeRequest) -> Result { + self.client .delete_range(req) .await .map_err(BoxedError::new) - .context(MetaSrvSnafu)?; - info!( - "Delete range, key: {}, end: {}, deleted: {}", - String::from_utf8_lossy(key), - String::from_utf8_lossy(end), - resp.deleted() - ); - - Ok(()) + .context(MetaSrvSnafu) } - async fn compare_and_set( + async fn batch_delete(&self, req: BatchDeleteRequest) -> Result { + self.client + .batch_delete(req) + .await + .map_err(BoxedError::new) + .context(MetaSrvSnafu) + } + + async fn batch_get(&self, req: BatchGetRequest) -> Result { + self.client + .batch_get(req) + .await + .map_err(BoxedError::new) + .context(MetaSrvSnafu) + } + + async fn compare_and_put( &self, - key: &[u8], - expect: &[u8], - val: &[u8], - ) -> Result>>> { - let request = CompareAndPutRequest::new() - .with_key(key.to_vec()) - .with_expect(expect.to_vec()) - .with_value(val.to_vec()); - let mut response = self - .client + request: CompareAndPutRequest, + ) -> Result { + self.client .compare_and_put(request) .await .map_err(BoxedError::new) - .context(MetaSrvSnafu)?; - if response.is_success() { - Ok(Ok(())) - } else { - Ok(Err(response.take_prev_kv().map(|v| v.value().to_vec()))) - } + .context(MetaSrvSnafu) } - async fn move_value(&self, from_key: &[u8], to_key: &[u8]) -> Result<()> { - let req = MoveValueRequest::new(from_key, to_key); - let _ = self - .client + async fn move_value(&self, req: MoveValueRequest) -> Result { + self.client .move_value(req) .await .map_err(BoxedError::new) - .context(MetaSrvSnafu)?; - Ok(()) + .context(MetaSrvSnafu) } fn as_any(&self) -> &dyn Any { diff --git a/src/catalog/src/remote/manager.rs b/src/catalog/src/remote/manager.rs index 0085aaf2ee..e826341af4 100644 --- a/src/catalog/src/remote/manager.rs +++ b/src/catalog/src/remote/manager.rs @@ -14,17 +14,15 @@ use std::any::Any; use std::collections::HashSet; -use std::pin::Pin; use std::sync::Arc; -use async_stream::stream; use async_trait::async_trait; use common_catalog::consts::{MAX_SYS_TABLE_ID, MITO_ENGINE}; use common_meta::ident::TableIdent; -use common_meta::kv_backend::{Kv, KvBackendRef}; +use common_meta::kv_backend::KvBackendRef; +use common_meta::rpc::store::{PutRequest, RangeRequest}; +use common_meta::rpc::KeyValue; use common_telemetry::{debug, error, info, warn}; -use futures::Stream; -use futures_util::{StreamExt, TryStreamExt}; use metrics::{decrement_gauge, increment_gauge}; use snafu::ResultExt; use table::engine::manager::TableEngineManagerRef; @@ -74,35 +72,39 @@ impl RemoteCatalogManager { } } - async fn iter_remote_catalogs( - &self, - ) -> Pin> + Send + '_>> { + async fn iter_remote_catalogs(&self) -> Result> { let catalog_range_prefix = build_catalog_prefix(); - let mut catalogs = self.backend.range(catalog_range_prefix.as_bytes()); - Box::pin(stream!({ - while let Some(r) = catalogs.next().await { - let Kv(k, _) = r.context(TableMetadataManagerSnafu)?; - if !k.starts_with(catalog_range_prefix.as_bytes()) { - debug!("Ignoring non-catalog key: {}", String::from_utf8_lossy(&k)); - continue; - } + let req = RangeRequest::new().with_prefix(catalog_range_prefix.as_bytes()); - let catalog_key = String::from_utf8_lossy(&k); - if let Ok(key) = CatalogKey::parse(&catalog_key) { - yield Ok(key) - } else { - error!("Invalid catalog key: {:?}", catalog_key); + let kvs = self + .backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; + + let catalogs = kvs + .into_iter() + .filter_map(|kv| { + let catalog_key = String::from_utf8_lossy(kv.key()); + + match CatalogKey::parse(&catalog_key) { + Ok(x) => Some(x), + Err(e) => { + error!(e; "Ignore invalid catalog key {:?}", catalog_key); + None + } } - } - })) + }) + .collect(); + Ok(catalogs) } /// Fetch catalogs/schemas/tables from remote catalog manager along with max table id allocated. async fn initiate_catalogs(&self) -> Result<()> { - let mut catalogs = self.iter_remote_catalogs().await; + let catalogs = self.iter_remote_catalogs().await?; let mut joins = Vec::new(); - while let Some(r) = catalogs.next().await { - let CatalogKey { catalog_name, .. } = r?; + for CatalogKey { catalog_name } in catalogs { info!("Fetch catalog from metasrv: {}", catalog_name); let node_id = self.node_id; @@ -128,13 +130,11 @@ impl RemoteCatalogManager { schema_name: schema_name.to_string(), } .to_string(); + let req = PutRequest::new() + .with_key(schema_key.as_bytes()) + .with_value(SchemaValue.as_bytes().context(InvalidCatalogValueSnafu)?); self.backend - .set( - schema_key.as_bytes(), - &SchemaValue {} - .as_bytes() - .context(InvalidCatalogValueSnafu)?, - ) + .put(req) .await .context(TableMetadataManagerSnafu)?; info!("Created schema '{schema_key}'"); @@ -143,13 +143,11 @@ impl RemoteCatalogManager { catalog_name: catalog_name.to_string(), } .to_string(); + let req = PutRequest::new() + .with_key(catalog_key.as_bytes()) + .with_value(CatalogValue.as_bytes().context(InvalidCatalogValueSnafu)?); self.backend - .set( - catalog_key.as_bytes(), - &CatalogValue {} - .as_bytes() - .context(InvalidCatalogValueSnafu)?, - ) + .put(req) .await .context(TableMetadataManagerSnafu)?; info!("Created catalog '{catalog_key}"); @@ -174,9 +172,7 @@ impl RemoteCatalogManager { schema_name: String, ) -> Result { info!("initializing tables in {}.{}", catalog_name, schema_name); - let tables = iter_remote_tables(node_id, &backend, &catalog_name, &schema_name).await; - - let kvs = tables.try_collect::>().await?; + let kvs = iter_remote_tables(node_id, &backend, &catalog_name, &schema_name).await?; let table_num = kvs.len(); let joins = kvs .into_iter() @@ -253,15 +249,14 @@ impl RemoteCatalogManager { engine_manager: TableEngineManagerRef, catalog_name: String, ) -> Result<()> { - let mut schemas = iter_remote_schemas(&backend, &catalog_name).await; - let mut joins = Vec::new(); - while let Some(r) = schemas.next().await { - let SchemaKey { - catalog_name, - schema_name, - .. - } = r?; + let schemas = iter_remote_schemas(&backend, &catalog_name).await?; + let mut joins = Vec::new(); + for SchemaKey { + catalog_name, + schema_name, + } in schemas + { info!( "Fetch schema from metasrv: {}.{}", &catalog_name, &schema_name @@ -314,11 +309,11 @@ impl RemoteCatalogManager { let table_key = self .build_regional_table_key(catalog_name, schema_name, table_name) .to_string(); + let req = PutRequest::new() + .with_key(table_key.as_bytes()) + .with_value(table_value.as_bytes().context(InvalidCatalogValueSnafu)?); self.backend - .set( - table_key.as_bytes(), - &table_value.as_bytes().context(InvalidCatalogValueSnafu)?, - ) + .put(req) .await .context(TableMetadataManagerSnafu)?; debug!( @@ -349,7 +344,7 @@ impl RemoteCatalogManager { .get(table_key.as_bytes()) .await .context(TableMetadataManagerSnafu)? - .map(|Kv(_, v)| { + .map(|KeyValue { key: _, value: v }| { let TableRegionalValue { table_id, engine_name, @@ -367,7 +362,7 @@ impl RemoteCatalogManager { }; self.backend - .delete(table_key.as_bytes()) + .delete(table_key.as_bytes(), false) .await .context(TableMetadataManagerSnafu)?; debug!( @@ -430,23 +425,30 @@ impl RemoteCatalogManager { async fn iter_remote_schemas<'a>( backend: &'a KvBackendRef, catalog_name: &'a str, -) -> Pin> + Send + 'a>> { +) -> Result> { let schema_prefix = build_schema_prefix(catalog_name); - let mut schemas = backend.range(schema_prefix.as_bytes()); + let req = RangeRequest::new().with_prefix(schema_prefix.as_bytes()); - Box::pin(stream!({ - while let Some(r) = schemas.next().await { - let Kv(k, _) = r.context(TableMetadataManagerSnafu)?; - if !k.starts_with(schema_prefix.as_bytes()) { - debug!("Ignoring non-schema key: {}", String::from_utf8_lossy(&k)); - continue; + let kvs = backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; + + let schemas = kvs + .into_iter() + .filter_map(|kv| { + let schema_key = String::from_utf8_lossy(kv.key()); + match SchemaKey::parse(&schema_key) { + Ok(x) => Some(x), + Err(e) => { + warn!("Ignore invalid schema key {:?}: {e}", schema_key); + None + } } - - let schema_key = - SchemaKey::parse(&String::from_utf8_lossy(&k)).context(InvalidCatalogValueSnafu)?; - yield Ok(schema_key) - } - })) + }) + .collect(); + Ok(schemas) } /// Iterate over all table entries on metasrv @@ -455,35 +457,42 @@ async fn iter_remote_tables<'a>( backend: &'a KvBackendRef, catalog_name: &'a str, schema_name: &'a str, -) -> Pin> + Send + 'a>> { +) -> Result> { let table_prefix = build_table_global_prefix(catalog_name, schema_name); - let mut tables = backend.range(table_prefix.as_bytes()); - Box::pin(stream!({ - while let Some(r) = tables.next().await { - let Kv(k, v) = r.context(TableMetadataManagerSnafu)?; - if !k.starts_with(table_prefix.as_bytes()) { - debug!("Ignoring non-table prefix: {}", String::from_utf8_lossy(&k)); - continue; - } - let table_key = TableGlobalKey::parse(&String::from_utf8_lossy(&k)) - .context(InvalidCatalogValueSnafu)?; - let table_value = TableGlobalValue::from_bytes(&v).context(InvalidCatalogValueSnafu)?; + let req = RangeRequest::new().with_prefix(table_prefix.as_bytes()); - info!( - "Found catalog table entry, key: {}, value: {:?}", - table_key, table_value - ); - // metasrv has allocated region ids to current datanode - if table_value - .regions_id_map - .get(&node_id) - .map(|v| !v.is_empty()) - .unwrap_or(false) - { - yield Ok((table_key, table_value)) - } + let kvs = backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; + + let mut tables = Vec::with_capacity(kvs.len()); + for kv in kvs { + let tgk = &String::from_utf8_lossy(kv.key()); + let Ok(table_key) = TableGlobalKey::parse(tgk) else { + warn!("Ignore invalid table global key {:?}", tgk); + continue; + }; + + let Ok(table_value) = TableGlobalValue::from_bytes(kv.value()) else { + warn!("Ignore invalid table global value {:?}", String::from_utf8_lossy(kv.value())); + continue; + }; + + info!("Found catalog table entry, key: {table_key}, value: {table_value:?}"); + + // metasrv has allocated region ids to current datanode + if table_value + .regions_id_map + .get(&node_id) + .map(|v| !v.is_empty()) + .unwrap_or(false) + { + tables.push((table_key, table_value)) } - })) + } + Ok(tables) } async fn print_regional_key_debug_info( @@ -500,7 +509,10 @@ async fn print_regional_key_debug_info( .to_string(); match backend.get(regional_key.as_bytes()).await { - Ok(Some(Kv(_, values_bytes))) => { + Ok(Some(KeyValue { + key: _, + value: values_bytes, + })) => { debug!( "Node id: {}, TableRegionalKey: {}, value: {},", node_id, @@ -702,13 +714,11 @@ impl CatalogManager for RemoteCatalogManager { let catalog_name = request.catalog; let schema_name = request.schema; let key = self.build_schema_key(catalog_name, schema_name).to_string(); + let req = PutRequest::new() + .with_key(key.as_bytes()) + .with_value(SchemaValue.as_bytes().context(InvalidCatalogValueSnafu)?); self.backend - .set( - key.as_bytes(), - &SchemaValue {} - .as_bytes() - .context(InvalidCatalogValueSnafu)?, - ) + .put(req) .await .context(TableMetadataManagerSnafu)?; @@ -729,21 +739,27 @@ impl CatalogManager for RemoteCatalogManager { node_id: self.node_id, } .to_string(); - let Some(Kv(_, value_bytes)) = self.backend.get(old_table_key.as_bytes()).await.context(TableMetadataManagerSnafu)? else { - return Ok(false) - }; + let Some(KeyValue{ key: _, value }) = self.backend + .get(old_table_key.as_bytes()) + .await + .context(TableMetadataManagerSnafu)? else { + return Ok(false) + }; let new_table_key = TableRegionalKey { catalog_name: request.catalog.clone(), schema_name: request.schema.clone(), table_name: request.new_table_name, node_id: self.node_id, }; + let req = PutRequest::new() + .with_key(new_table_key.to_string().as_bytes()) + .with_value(value); self.backend - .set(new_table_key.to_string().as_bytes(), &value_bytes) + .put(req) .await .context(TableMetadataManagerSnafu)?; self.backend - .delete(old_table_key.to_string().as_bytes()) + .delete(old_table_key.to_string().as_bytes(), false) .await .context(TableMetadataManagerSnafu)?; Ok(true) @@ -796,7 +812,7 @@ impl CatalogManager for RemoteCatalogManager { .get(key.as_bytes()) .await .context(TableMetadataManagerSnafu)? - .map(|Kv(_, v)| { + .map(|KeyValue { key: _, value: v }| { let TableRegionalValue { table_id, engine_name, @@ -863,16 +879,19 @@ impl CatalogManager for RemoteCatalogManager { } async fn catalog_names(&self) -> Result> { - let mut stream = self.backend.range(CATALOG_KEY_PREFIX.as_bytes()); + let req = RangeRequest::new().with_prefix(CATALOG_KEY_PREFIX.as_bytes()); + let kvs = self + .backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; let mut catalogs = HashSet::new(); - while let Some(catalog) = stream.next().await { - if let Ok(catalog) = catalog { - let catalog_key = String::from_utf8_lossy(&catalog.0); - - if let Ok(key) = CatalogKey::parse(&catalog_key) { - let _ = catalogs.insert(key.catalog_name); - } + for catalog in kvs { + let catalog_key = String::from_utf8_lossy(catalog.key()); + if let Ok(key) = CatalogKey::parse(&catalog_key) { + let _ = catalogs.insert(key.catalog_name); } } @@ -880,18 +899,19 @@ impl CatalogManager for RemoteCatalogManager { } async fn schema_names(&self, catalog_name: &str) -> Result> { - let mut stream = self + let req = RangeRequest::new().with_prefix(build_schema_prefix(catalog_name).as_bytes()); + let kvs = self .backend - .range(build_schema_prefix(catalog_name).as_bytes()); + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; let mut schemas = HashSet::new(); - while let Some(schema) = stream.next().await { - if let Ok(schema) = schema { - let schema_key = String::from_utf8_lossy(&schema.0); - - if let Ok(key) = SchemaKey::parse(&schema_key) { - let _ = schemas.insert(key.schema_name); - } + for schema in kvs { + let schema_key = String::from_utf8_lossy(schema.key()); + if let Ok(key) = SchemaKey::parse(&schema_key) { + let _ = schemas.insert(key.schema_name); } } Ok(schemas.into_iter().collect()) @@ -901,18 +921,20 @@ impl CatalogManager for RemoteCatalogManager { self.check_catalog_schema_exist(catalog_name, schema_name) .await?; - let mut stream = self + let req = RangeRequest::new() + .with_prefix(build_table_regional_prefix(catalog_name, schema_name).as_bytes()); + let kvs = self .backend - .range(build_table_regional_prefix(catalog_name, schema_name).as_bytes()); + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; let mut tables = HashSet::new(); - while let Some(table) = stream.next().await { - if let Ok(table) = table { - let table_key = String::from_utf8_lossy(&table.0); - - if let Ok(key) = TableRegionalKey::parse(&table_key) { - let _ = tables.insert(key.table_name); - } + for table in kvs { + let table_key = String::from_utf8_lossy(table.key()); + if let Ok(key) = TableRegionalKey::parse(&table_key) { + let _ = tables.insert(key.table_name); } } Ok(tables.into_iter().collect()) @@ -921,13 +943,11 @@ impl CatalogManager for RemoteCatalogManager { async fn register_catalog(&self, name: String) -> Result { let key = CatalogKey { catalog_name: name }.to_string(); // TODO(hl): use compare_and_swap to prevent concurrent update + let req = PutRequest::new() + .with_key(key.as_bytes()) + .with_value(CatalogValue.as_bytes().context(InvalidCatalogValueSnafu)?); self.backend - .set( - key.as_bytes(), - &CatalogValue {} - .as_bytes() - .context(InvalidCatalogValueSnafu)?, - ) + .put(req) .await .context(TableMetadataManagerSnafu)?; increment_gauge!(crate::metrics::METRIC_CATALOG_MANAGER_CATALOG_COUNT, 1.0); diff --git a/src/catalog/tests/remote_catalog_tests.rs b/src/catalog/tests/remote_catalog_tests.rs index fb0d3fd598..f43011b60e 100644 --- a/src/catalog/tests/remote_catalog_tests.rs +++ b/src/catalog/tests/remote_catalog_tests.rs @@ -21,6 +21,7 @@ mod tests { use std::sync::Arc; use std::time::Duration; + use catalog::error::Error; use catalog::helper::{CatalogKey, CatalogValue, SchemaKey, SchemaValue}; use catalog::remote::mock::MockTableEngine; use catalog::remote::region_alive_keeper::RegionAliveKeepers; @@ -30,8 +31,8 @@ mod tests { use common_meta::ident::TableIdent; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::KvBackend; + use common_meta::rpc::store::{CompareAndPutRequest, PutRequest, RangeRequest}; use datatypes::schema::RawSchema; - use futures_util::StreamExt; use table::engine::manager::{MemoryTableEngineManager, TableEngineManagerRef}; use table::engine::{EngineContext, TableEngineRef}; use table::requests::CreateTableRequest; @@ -52,38 +53,35 @@ mod tests { #[tokio::test] async fn test_backend() { - common_telemetry::init_default_ut_logging(); - let backend = MemoryKvBackend::default(); + let backend = MemoryKvBackend::::default(); let default_catalog_key = CatalogKey { catalog_name: DEFAULT_CATALOG_NAME.to_string(), } .to_string(); - - backend - .set( - default_catalog_key.as_bytes(), - &CatalogValue {}.as_bytes().unwrap(), - ) - .await - .unwrap(); + let req = PutRequest::new() + .with_key(default_catalog_key.as_bytes()) + .with_value(CatalogValue.as_bytes().unwrap()); + backend.put(req).await.unwrap(); let schema_key = SchemaKey { catalog_name: DEFAULT_CATALOG_NAME.to_string(), schema_name: DEFAULT_SCHEMA_NAME.to_string(), } .to_string(); - backend - .set(schema_key.as_bytes(), &SchemaValue {}.as_bytes().unwrap()) - .await - .unwrap(); + let req = PutRequest::new() + .with_key(schema_key.as_bytes()) + .with_value(SchemaValue.as_bytes().unwrap()); + backend.put(req).await.unwrap(); - let mut iter = backend.range("__c-".as_bytes()); - let mut res = HashSet::new(); - while let Some(r) = iter.next().await { - let kv = r.unwrap(); - let _ = res.insert(String::from_utf8_lossy(&kv.0).to_string()); - } + let req = RangeRequest::new().with_prefix(b"__c-".to_vec()); + let res = backend + .range(req) + .await + .unwrap() + .kvs + .into_iter() + .map(|kv| String::from_utf8_lossy(kv.key()).to_string()); assert_eq!( vec!["__c-greptime".to_string()], res.into_iter().collect::>() @@ -98,36 +96,32 @@ mod tests { catalog_name: DEFAULT_CATALOG_NAME.to_string(), } .to_string(); - - backend - .set( - default_catalog_key.as_bytes(), - &CatalogValue {}.as_bytes().unwrap(), - ) - .await - .unwrap(); + let req = PutRequest::new() + .with_key(default_catalog_key.as_bytes()) + .with_value(CatalogValue.as_bytes().unwrap()); + backend.put(req).await.unwrap(); let ret = backend.get(b"__c-greptime").await.unwrap(); let _ = ret.unwrap(); - let _ = backend - .compare_and_set( - b"__c-greptime", - &CatalogValue {}.as_bytes().unwrap(), - b"123", - ) - .await - .unwrap(); + let req = CompareAndPutRequest::new() + .with_key(b"__c-greptime".to_vec()) + .with_expect(CatalogValue.as_bytes().unwrap()) + .with_value(b"123".to_vec()); + let _ = backend.compare_and_put(req).await.unwrap(); let ret = backend.get(b"__c-greptime").await.unwrap(); - assert_eq!(&b"123"[..], &(ret.as_ref().unwrap().1)); + assert_eq!(b"123", ret.as_ref().unwrap().value.as_slice()); - let _ = backend.set(b"__c-greptime", b"1234").await; + let req = PutRequest::new() + .with_key(b"__c-greptime".to_vec()) + .with_value(b"1234".to_vec()); + let _ = backend.put(req).await; let ret = backend.get(b"__c-greptime").await.unwrap(); - assert_eq!(&b"1234"[..], &(ret.as_ref().unwrap().1)); + assert_eq!(b"1234", ret.unwrap().value.as_slice()); - backend.delete(b"__c-greptime").await.unwrap(); + backend.delete(b"__c-greptime", false).await.unwrap(); let ret = backend.get(b"__c-greptime").await.unwrap(); assert!(ret.is_none()); @@ -135,8 +129,16 @@ mod tests { async fn prepare_components(node_id: u64) -> TestingComponents { let backend = Arc::new(MemoryKvBackend::default()); - backend.set(b"__c-greptime", b"").await.unwrap(); - backend.set(b"__s-greptime-public", b"").await.unwrap(); + + let req = PutRequest::new() + .with_key(b"__c-greptime".to_vec()) + .with_value(b"".to_vec()); + backend.put(req).await.unwrap(); + + let req = PutRequest::new() + .with_key(b"__s-greptime-public".to_vec()) + .with_value(b"".to_vec()); + backend.put(req).await.unwrap(); let cached_backend = Arc::new(CachedMetaKvBackend::wrap(backend)); diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 355d17e51a..027fe0e047 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -13,6 +13,7 @@ common-error = { path = "../error" } common-runtime = { path = "../runtime" } common-telemetry = { path = "../telemetry" } common-time = { path = "../time" } +etcd-client.workspace = true futures.workspace = true lazy_static.workspace = true prost.workspace = true diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 9d541705b9..b1a7cbd5d5 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -15,6 +15,8 @@ use common_error::prelude::*; use serde_json::error::Error as JsonError; use snafu::Location; +use store_api::storage::RegionNumber; +use table::metadata::TableId; #[derive(Debug, Snafu)] #[snafu(visibility(pub))] @@ -73,6 +75,22 @@ pub enum Error { source: BoxedError, location: Location, }, + + #[snafu(display("Etcd txn error: {err_msg}"))] + EtcdTxnOpResponse { err_msg: String, location: Location }, + + #[snafu(display( + "Failed to move region {} in table {}, err: {}", + region, + table_id, + err_msg + ))] + MoveRegion { + table_id: TableId, + region: RegionNumber, + err_msg: String, + location: Location, + }, } pub type Result = std::result::Result; @@ -81,12 +99,13 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { - IllegalServerState { .. } => StatusCode::Internal, + IllegalServerState { .. } | EtcdTxnOpResponse { .. } => StatusCode::Internal, SerdeJson { .. } | RouteInfoCorrupted { .. } | InvalidProtoMsg { .. } - | InvalidTableMetadata { .. } => StatusCode::Unexpected, + | InvalidTableMetadata { .. } + | MoveRegion { .. } => StatusCode::Unexpected, SendMessage { .. } | GetKvCache { .. } diff --git a/src/common/meta/src/key/datanode_table.rs b/src/common/meta/src/key/datanode_table.rs index 04bc2867b4..58cfb9de66 100644 --- a/src/common/meta/src/key/datanode_table.rs +++ b/src/common/meta/src/key/datanode_table.rs @@ -12,16 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -use futures::{StreamExt, TryStreamExt}; use serde::{Deserialize, Serialize}; -use snafu::OptionExt; +use snafu::{ensure, OptionExt}; use store_api::storage::RegionNumber; use table::metadata::TableId; use super::{DATANODE_TABLE_KEY_PATTERN, DATANODE_TABLE_KEY_PREFIX}; -use crate::error::{ConcurrentModifyRegionsPlacementSnafu, InvalidTableMetadataSnafu, Result}; +use crate::error::{ + ConcurrentModifyRegionsPlacementSnafu, InvalidTableMetadataSnafu, MoveRegionSnafu, Result, +}; use crate::key::{to_removed_key, TableMetaKey}; +use crate::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp}; use crate::kv_backend::KvBackendRef; +use crate::rpc::store::{CompareAndPutRequest, MoveValueRequest, RangeRequest}; use crate::DatanodeId; struct DatanodeTableKey { @@ -100,7 +103,7 @@ impl DatanodeTableManager { self.kv_backend .get(&key.as_raw_key()) .await? - .map(|kv| DatanodeTableValue::try_from_raw_value(kv.1)) + .map(|kv| DatanodeTableValue::try_from_raw_value(kv.value)) .transpose() } @@ -112,72 +115,123 @@ impl DatanodeTableManager { ) -> Result<()> { let key = DatanodeTableKey::new(datanode_id, table_id).as_raw_key(); let val = DatanodeTableValue::new(table_id, regions).try_as_raw_value()?; - self.kv_backend - .compare_and_set(&key, &[], &val) - .await? - .map_err(|curr| { - let curr = if let Some(curr) = curr { - DatanodeTableValue::try_from_raw_value(curr).map_or_else( - |e| format!("invalid DatanodeTableValue for Datanode {datanode_id}: {e}"), + let req = CompareAndPutRequest::new().with_key(key).with_value(val); + + let resp = self.kv_backend.compare_and_put(req).await?; + if !resp.success { + let curr = resp.prev_kv.map_or_else( + || "empty".to_string(), + |kv| { + DatanodeTableValue::try_from_raw_value(kv.value).map_or_else( + |e| format!("Invalid DatanodeTableValue for Datanode {datanode_id}: {e}"), |v| format!("{v:?}"), ) - } else { - "empty".to_string() - }; - ConcurrentModifyRegionsPlacementSnafu { - err_msg: format!("Datanode {datanode_id} already existed {curr}"), - } - .build() - }) + }, + ); + return ConcurrentModifyRegionsPlacementSnafu { + err_msg: format!("Datanode {datanode_id} already existed {curr}"), + } + .fail(); + } + Ok(()) } pub async fn remove(&self, datanode_id: DatanodeId, table_id: TableId) -> Result<()> { let key = DatanodeTableKey::new(datanode_id, table_id); let removed_key = to_removed_key(&String::from_utf8_lossy(&key.as_raw_key())); - self.kv_backend - .move_value(&key.as_raw_key(), removed_key.as_bytes()) - .await + let req = MoveValueRequest::new(key.as_raw_key(), removed_key.as_bytes()); + let _ = self.kv_backend.move_value(req).await?; + Ok(()) } - // TODO(LFC): Use transaction to move region once the KvBackend and KvStore are merged into one. pub async fn move_region( &self, from_datanode: DatanodeId, to_datanode: DatanodeId, table_id: TableId, region: RegionNumber, - ) -> Result<()> { + ) -> Result { let from_key = DatanodeTableKey::new(from_datanode, table_id); - let from_value = self.get(&from_key).await?; - if let Some(mut from_value) = from_value { - from_value.regions.retain(|x| *x != region); - from_value.version += 1; - self.kv_backend - .set(&from_key.as_raw_key(), &from_value.try_as_raw_value()?) - .await?; - } + let mut from_value = self.get(&from_key).await?.context(MoveRegionSnafu { + table_id, + region, + err_msg: format!("DatanodeTableKey not found for Datanode {from_datanode}"), + })?; + + ensure!( + from_value.regions.contains(®ion), + MoveRegionSnafu { + table_id, + region, + err_msg: format!("target region not found in Datanode {from_datanode}"), + } + ); let to_key = DatanodeTableKey::new(to_datanode, table_id); let to_value = self.get(&to_key).await?; - if let Some(mut to_value) = to_value { - to_value.regions.push(region); - to_value.version += 1; - self.kv_backend - .set(&to_key.as_raw_key(), &to_value.try_as_raw_value()?) - .await?; + + if let Some(v) = to_value.as_ref() { + ensure!( + !v.regions.contains(®ion), + MoveRegionSnafu { + table_id, + region, + err_msg: format!("target region already existed in Datanode {to_datanode}"), + } + ); } - Ok(()) + + let compares = vec![ + Compare::with_value( + from_key.as_raw_key(), + CompareOp::Equal, + from_value.try_as_raw_value()?, + ), + Compare::new( + to_key.as_raw_key(), + CompareOp::Equal, + to_value + .as_ref() + .map(|x| x.try_as_raw_value()) + .transpose()?, + ), + ]; + + let mut operations = Vec::with_capacity(2); + + from_value.regions.retain(|x| *x != region); + if from_value.regions.is_empty() { + operations.push(TxnOp::Delete(from_key.as_raw_key())); + } else { + from_value.version += 1; + operations.push(TxnOp::Put( + from_key.as_raw_key(), + from_value.try_as_raw_value()?, + )); + } + + if let Some(mut v) = to_value { + v.regions.push(region); + v.version += 1; + operations.push(TxnOp::Put(to_key.as_raw_key(), v.try_as_raw_value()?)); + } else { + let v = DatanodeTableValue::new(table_id, vec![region]); + operations.push(TxnOp::Put(to_key.as_raw_key(), v.try_as_raw_value()?)); + } + + let txn = Txn::new().when(compares).and_then(operations); + let resp = self.kv_backend.txn(txn).await?; + Ok(resp.succeeded) } pub async fn tables(&self, datanode_id: DatanodeId) -> Result> { let prefix = DatanodeTableKey::prefix(datanode_id); - let table_ids = self - .kv_backend - .range(prefix.as_bytes()) - .map(|result| result.map(|kv| DatanodeTableValue::try_from_raw_value(kv.1))) - .try_collect::>() - .await? + let req = RangeRequest::new().with_prefix(prefix.as_bytes()); + let resp = self.kv_backend.range(req).await?; + let table_ids = resp + .kvs .into_iter() + .map(|kv| DatanodeTableValue::try_from_raw_value(kv.value)) .collect::>>()?; Ok(table_ids) } @@ -195,11 +249,20 @@ mod tests { async fn test_move_region() { let manager = DatanodeTableManager::new(Arc::new(MemoryKvBackend::default())); - assert!(manager.create(1, 1, vec![1, 2]).await.is_ok()); - assert!(manager.create(2, 1, vec![3, 4]).await.is_ok()); + let result = manager.move_region(1, 2, 1, 1).await; + assert!(result.unwrap_err().to_string().contains( + "Failed to move region 1 in table 1, err: DatanodeTableKey not found for Datanode 1" + )); - assert!(manager.move_region(1, 2, 1, 1).await.is_ok()); + assert!(manager.create(1, 1, vec![1, 2, 3]).await.is_ok()); + let result = manager.move_region(1, 2, 1, 100).await; + assert!(result.unwrap_err().to_string().contains( + "Failed to move region 100 in table 1, err: target region not found in Datanode 1" + )); + // Move region 1 from datanode 1 to datanode 2. + // Note that the DatanodeTableValue is not existed for datanode 2 now. + assert!(manager.move_region(1, 2, 1, 1).await.unwrap()); let value = manager .get(&DatanodeTableKey::new(1, 1)) .await @@ -209,11 +272,10 @@ mod tests { value, DatanodeTableValue { table_id: 1, - regions: vec![2], + regions: vec![2, 3], version: 1, } ); - let value = manager .get(&DatanodeTableKey::new(2, 1)) .await @@ -223,10 +285,57 @@ mod tests { value, DatanodeTableValue { table_id: 1, - regions: vec![3, 4, 1], + regions: vec![1], + version: 0, + } + ); + + // Move region 2 from datanode 1 to datanode 2. + assert!(manager.move_region(1, 2, 1, 2).await.is_ok()); + let value = manager + .get(&DatanodeTableKey::new(1, 1)) + .await + .unwrap() + .unwrap(); + assert_eq!( + value, + DatanodeTableValue { + table_id: 1, + regions: vec![3], + version: 2, + } + ); + let value = manager + .get(&DatanodeTableKey::new(2, 1)) + .await + .unwrap() + .unwrap(); + assert_eq!( + value, + DatanodeTableValue { + table_id: 1, + regions: vec![1, 2], version: 1, } ); + + // Move region 3 (the last region) from datanode 1 to datanode 2. + assert!(manager.move_region(1, 2, 1, 3).await.is_ok()); + let value = manager.get(&DatanodeTableKey::new(1, 1)).await.unwrap(); + assert!(value.is_none()); + let value = manager + .get(&DatanodeTableKey::new(2, 1)) + .await + .unwrap() + .unwrap(); + assert_eq!( + value, + DatanodeTableValue { + table_id: 1, + regions: vec![1, 2, 3], + version: 2, + } + ); } #[tokio::test] @@ -262,8 +371,8 @@ mod tests { .await .unwrap() .unwrap(); - assert_eq!(b"__removed-__dn_table/2/1", kv.0.as_slice()); - let value = DatanodeTableValue::try_from_raw_value(kv.1).unwrap(); + assert_eq!(b"__removed-__dn_table/2/1", kv.key()); + let value = DatanodeTableValue::try_from_raw_value(kv.value).unwrap(); assert_eq!(value, expected_value); let values = manager.tables(1).await.unwrap(); diff --git a/src/common/meta/src/key/table_info.rs b/src/common/meta/src/key/table_info.rs index 1920863436..9f3e9974b8 100644 --- a/src/common/meta/src/key/table_info.rs +++ b/src/common/meta/src/key/table_info.rs @@ -19,6 +19,7 @@ use super::TABLE_INFO_KEY_PREFIX; use crate::error::Result; use crate::key::{to_removed_key, TableMetaKey}; use crate::kv_backend::KvBackendRef; +use crate::rpc::store::{CompareAndPutRequest, MoveValueRequest}; pub struct TableInfoKey { table_id: TableId, @@ -57,11 +58,18 @@ impl TableInfoManager { self.kv_backend .get(&raw_key) .await? - .map(|x| TableInfoValue::try_from_raw_value(x.1)) + .map(|x| TableInfoValue::try_from_raw_value(x.value)) .transpose() } - pub async fn compare_and_set( + /// Compare and put value of key. `expect` is the expected value, if backend's current value associated + /// with key is the same as `expect`, the value will be updated to `val`. + /// + /// - If the compare-and-set operation successfully updated value, this method will return an `Ok(Ok())` + /// - If associated value is not the same as `expect`, no value will be updated and an `Ok(Err(Vec))` + /// will be returned, the `Err(Vec)` indicates the current associated value of key. + /// - If any error happens during operation, an `Err(Error)` will be returned. + pub async fn compare_and_put( &self, table_id: TableId, expect: Option, @@ -82,17 +90,24 @@ impl TableInfoManager { }; let raw_value = value.try_as_raw_value()?; - self.kv_backend - .compare_and_set(&raw_key, &expect, &raw_value) - .await + let req = CompareAndPutRequest::new() + .with_key(raw_key) + .with_expect(expect) + .with_value(raw_value); + let resp = self.kv_backend.compare_and_put(req).await?; + Ok(if resp.success { + Ok(()) + } else { + Err(resp.prev_kv.map(|x| x.value)) + }) } pub async fn remove(&self, table_id: TableId) -> Result<()> { - let key = TableInfoKey::new(table_id); - let removed_key = to_removed_key(&String::from_utf8_lossy(key.as_raw_key().as_slice())); - self.kv_backend - .move_value(&key.as_raw_key(), removed_key.as_bytes()) - .await + let key = TableInfoKey::new(table_id).as_raw_key(); + let removed_key = to_removed_key(&String::from_utf8_lossy(&key)); + let req = MoveValueRequest::new(key, removed_key.as_bytes()); + self.kv_backend.move_value(req).await?; + Ok(()) } } @@ -107,6 +122,7 @@ mod tests { use super::*; use crate::kv_backend::memory::MemoryKvBackend; use crate::kv_backend::KvBackend; + use crate::rpc::store::PutRequest; #[tokio::test] async fn test_table_info_manager() { @@ -120,7 +136,8 @@ mod tests { } .try_as_raw_value() .unwrap(); - backend.set(&key, &val).await.unwrap(); + let req = PutRequest::new().with_key(key).with_value(val); + backend.put(req).await.unwrap(); } let manager = TableInfoManager::new(backend.clone()); @@ -137,7 +154,7 @@ mod tests { let table_info = new_table_info(4); let result = manager - .compare_and_set(4, None, table_info.clone()) + .compare_and_put(4, None, table_info.clone()) .await .unwrap(); assert!(result.is_ok()); @@ -145,7 +162,7 @@ mod tests { // test cas failed, the new table info is not set let new_table_info = new_table_info(4); let result = manager - .compare_and_set(4, None, new_table_info.clone()) + .compare_and_put(4, None, new_table_info.clone()) .await .unwrap(); let actual = TableInfoValue::try_from_raw_value(result.unwrap_err().unwrap()).unwrap(); @@ -159,7 +176,7 @@ mod tests { // test cas success let result = manager - .compare_and_set(4, Some(actual), new_table_info.clone()) + .compare_and_put(4, Some(actual), new_table_info.clone()) .await .unwrap(); assert!(result.is_ok()); @@ -171,8 +188,8 @@ mod tests { .await .unwrap() .unwrap(); - assert_eq!(b"__removed-__table_info/4", kv.0.as_slice()); - let value = TableInfoValue::try_from_raw_value(kv.1).unwrap(); + assert_eq!(b"__removed-__table_info/4", kv.key.as_slice()); + let value = TableInfoValue::try_from_raw_value(kv.value).unwrap(); assert_eq!(value.table_info, new_table_info); assert_eq!(value.version, 1); } diff --git a/src/common/meta/src/key/table_name.rs b/src/common/meta/src/key/table_name.rs index d7cfcb2308..c8026d1761 100644 --- a/src/common/meta/src/key/table_name.rs +++ b/src/common/meta/src/key/table_name.rs @@ -14,7 +14,6 @@ use std::sync::Arc; -use futures::{StreamExt, TryStreamExt}; use serde::{Deserialize, Serialize}; use snafu::OptionExt; use table::metadata::TableId; @@ -24,6 +23,7 @@ use crate::error::{InvalidTableMetadataSnafu, Result}; use crate::key::{to_removed_key, TableMetaKey}; use crate::kv_backend::memory::MemoryKvBackend; use crate::kv_backend::KvBackendRef; +use crate::rpc::store::{CompareAndPutRequest, MoveValueRequest, RangeRequest}; #[derive(Debug)] pub struct TableNameKey<'a> { @@ -111,11 +111,11 @@ impl TableNameManager { let raw_key = key.as_raw_key(); let value = TableNameValue::new(table_id); let raw_value = value.try_as_raw_value()?; - let result = self - .kv_backend - .compare_and_set(&raw_key, &[], &raw_value) - .await?; - Ok(matches!(result, Ok(()))) + let req = CompareAndPutRequest::new() + .with_key(raw_key) + .with_value(raw_value); + let result = self.kv_backend.compare_and_put(req).await?; + Ok(result.success) } pub async fn get(&self, key: &TableNameKey<'_>) -> Result> { @@ -123,19 +123,18 @@ impl TableNameManager { self.kv_backend .get(&raw_key) .await? - .map(|x| TableNameValue::try_from_raw_value(x.1)) + .map(|x| TableNameValue::try_from_raw_value(x.value)) .transpose() } pub async fn tables(&self, catalog: &str, schema: &str) -> Result> { let key = TableNameKey::prefix_to_table(catalog, schema).into_bytes(); - let table_names = self - .kv_backend - .range(&key) - .map(|x| x.map(|kv| TableNameKey::strip_table_name(&kv.0))) - .try_collect::>() - .await? + let req = RangeRequest::new().with_prefix(key); + let resp = self.kv_backend.range(req).await?; + let table_names = resp + .kvs .into_iter() + .map(|kv| TableNameKey::strip_table_name(kv.key())) .collect::>>()?; Ok(table_names) } @@ -143,9 +142,9 @@ impl TableNameManager { pub async fn remove(&self, key: &TableNameKey<'_>) -> Result<()> { let raw_key = key.as_raw_key(); let removed_key = to_removed_key(&String::from_utf8_lossy(&raw_key)); - self.kv_backend - .move_value(&raw_key, removed_key.as_bytes()) - .await + let req = MoveValueRequest::new(raw_key, removed_key.as_bytes()); + let _ = self.kv_backend.move_value(req).await?; + Ok(()) } } @@ -183,7 +182,7 @@ mod tests { .await .unwrap() .unwrap(); - let value = TableNameValue::try_from_raw_value(kv.1).unwrap(); + let value = TableNameValue::try_from_raw_value(kv.value).unwrap(); assert_eq!(value.table_id(), 99); let tables = manager.tables("my_catalog", "my_schema").await.unwrap(); diff --git a/src/common/meta/src/key/table_region.rs b/src/common/meta/src/key/table_region.rs index 7d6254d3fb..50c17a2443 100644 --- a/src/common/meta/src/key/table_region.rs +++ b/src/common/meta/src/key/table_region.rs @@ -22,6 +22,7 @@ use super::TABLE_REGION_KEY_PREFIX; use crate::error::Result; use crate::key::{to_removed_key, TableMetaKey}; use crate::kv_backend::KvBackendRef; +use crate::rpc::store::{CompareAndPutRequest, MoveValueRequest}; use crate::DatanodeId; pub type RegionDistribution = BTreeMap>; @@ -63,11 +64,18 @@ impl TableRegionManager { self.kv_backend .get(&raw_key) .await? - .map(|x| TableRegionValue::try_from_raw_value(x.1)) + .map(|x| TableRegionValue::try_from_raw_value(x.value)) .transpose() } - pub async fn compare_and_set( + /// Compare and put value of key. `expect` is the expected value, if backend's current value associated + /// with key is the same as `expect`, the value will be updated to `val`. + /// + /// - If the compare-and-set operation successfully updated value, this method will return an `Ok(Ok())` + /// - If associated value is not the same as `expect`, no value will be updated and an `Ok(Err(Vec))` + /// will be returned, the `Err(Vec)` indicates the current associated value of key. + /// - If any error happens during operation, an `Err(Error)` will be returned. + pub async fn compare_and_put( &self, table_id: TableId, expect: Option, @@ -88,17 +96,24 @@ impl TableRegionManager { }; let raw_value = value.try_as_raw_value()?; - self.kv_backend - .compare_and_set(&raw_key, &expect, &raw_value) - .await + let req = CompareAndPutRequest::new() + .with_key(raw_key) + .with_expect(expect) + .with_value(raw_value); + let resp = self.kv_backend.compare_and_put(req).await?; + Ok(if resp.success { + Ok(()) + } else { + Err(resp.prev_kv.map(|x| x.value)) + }) } pub async fn remove(&self, table_id: TableId) -> Result<()> { - let key = TableRegionKey::new(table_id); - let remove_key = to_removed_key(&String::from_utf8_lossy(key.as_raw_key().as_slice())); - self.kv_backend - .move_value(&key.as_raw_key(), remove_key.as_bytes()) - .await + let key = TableRegionKey::new(table_id).as_raw_key(); + let remove_key = to_removed_key(&String::from_utf8_lossy(&key)); + let req = MoveValueRequest::new(key, remove_key.as_bytes()); + self.kv_backend.move_value(req).await?; + Ok(()) } } @@ -118,7 +133,7 @@ mod tests { let region_distribution = RegionDistribution::from([(1, vec![1, 2, 3]), (2, vec![4, 5, 6])]); let result = manager - .compare_and_set(1, None, region_distribution.clone()) + .compare_and_put(1, None, region_distribution.clone()) .await .unwrap(); assert!(result.is_ok()); @@ -126,7 +141,7 @@ mod tests { let new_region_distribution = RegionDistribution::from([(1, vec![4, 5, 6]), (2, vec![1, 2, 3])]); let curr = manager - .compare_and_set(1, None, new_region_distribution.clone()) + .compare_and_put(1, None, new_region_distribution.clone()) .await .unwrap() .unwrap_err() @@ -141,7 +156,7 @@ mod tests { ); assert!(manager - .compare_and_set(1, Some(curr), new_region_distribution.clone()) + .compare_and_put(1, Some(curr), new_region_distribution.clone()) .await .unwrap() .is_ok()); @@ -163,8 +178,8 @@ mod tests { .await .unwrap() .unwrap(); - assert_eq!(b"__removed-__table_region/1", kv.0.as_slice()); - let value = TableRegionValue::try_from_raw_value(kv.1).unwrap(); + assert_eq!(b"__removed-__table_region/1", kv.key.as_slice()); + let value = TableRegionValue::try_from_raw_value(kv.value).unwrap(); assert_eq!(value.region_distribution, new_region_distribution); assert_eq!(value.version, 1); } diff --git a/src/common/meta/src/kv_backend.rs b/src/common/meta/src/kv_backend.rs index 81ff2a3cb3..d068782def 100644 --- a/src/common/meta/src/kv_backend.rs +++ b/src/common/meta/src/kv_backend.rs @@ -13,68 +13,84 @@ // limitations under the License. pub mod memory; +pub mod txn; use std::any::Any; -use std::pin::Pin; use std::sync::Arc; use async_trait::async_trait; use common_error::ext::ErrorExt; -use futures::{Stream, StreamExt}; +pub use txn::TxnService; use crate::error::Error; +use crate::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, + BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, + DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, + RangeRequest, RangeResponse, +}; +use crate::rpc::KeyValue; -#[derive(Debug, Clone, PartialEq)] -pub struct Kv(pub Vec, pub Vec); - -pub type ValueIter<'a, E> = Pin> + Send + 'a>>; - -pub type KvBackendRef = Arc>; +pub type KvBackendRef = Arc + Send + Sync>; #[async_trait] -pub trait KvBackend: Send + Sync { - type Error: ErrorExt; +pub trait KvBackend: TxnService +where + Self::Error: ErrorExt, +{ + fn name(&self) -> &str; - fn range<'a, 'b>(&'a self, key: &[u8]) -> ValueIter<'b, Self::Error> - where - 'a: 'b; + async fn range(&self, req: RangeRequest) -> Result; - async fn set(&self, key: &[u8], val: &[u8]) -> Result<(), Self::Error>; + async fn put(&self, req: PutRequest) -> Result; - /// Compare and set value of key. `expect` is the expected value, if backend's current value associated - /// with key is the same as `expect`, the value will be updated to `val`. - /// - /// - If the compare-and-set operation successfully updated value, this method will return an `Ok(Ok())` - /// - If associated value is not the same as `expect`, no value will be updated and an `Ok(Err(Vec))` - /// will be returned, the `Err(Vec)` indicates the current associated value of key. - /// - If any error happens during operation, an `Err(Error)` will be returned. - async fn compare_and_set( + async fn batch_put(&self, req: BatchPutRequest) -> Result; + + async fn compare_and_put( &self, - key: &[u8], - expect: &[u8], - val: &[u8], - ) -> Result>>, Self::Error>; + req: CompareAndPutRequest, + ) -> Result; - async fn delete_range(&self, key: &[u8], end: &[u8]) -> Result<(), Self::Error>; + async fn delete_range( + &self, + req: DeleteRangeRequest, + ) -> Result; - async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> { - self.delete_range(key, &[]).await + async fn delete(&self, key: &[u8], prev_kv: bool) -> Result, Self::Error> { + let mut req = DeleteRangeRequest::new().with_key(key.to_vec()); + if prev_kv { + req = req.with_prev_kv(); + } + + let resp = self.delete_range(req).await?; + + if prev_kv { + Ok(resp.prev_kvs.into_iter().next()) + } else { + Ok(None) + } } + async fn batch_delete( + &self, + req: BatchDeleteRequest, + ) -> Result; + /// Default get is implemented based on `range` method. - async fn get(&self, key: &[u8]) -> Result, Self::Error> { - let mut iter = self.range(key); - while let Some(r) = iter.next().await { - let kv = r?; - if kv.0 == key { - return Ok(Some(kv)); - } - } - return Ok(None); + async fn get(&self, key: &[u8]) -> Result, Self::Error> { + let req = RangeRequest::new().with_key(key.to_vec()); + let mut resp = self.range(req).await?; + Ok(if resp.kvs.is_empty() { + None + } else { + Some(resp.kvs.remove(0)) + }) } + async fn batch_get(&self, req: BatchGetRequest) -> Result; + /// MoveValue atomically renames the key to the given updated key. - async fn move_value(&self, from_key: &[u8], to_key: &[u8]) -> Result<(), Self::Error>; + async fn move_value(&self, req: MoveValueRequest) -> Result; fn as_any(&self) -> &dyn Any; } diff --git a/src/common/meta/src/kv_backend/memory.rs b/src/common/meta/src/kv_backend/memory.rs index 6108120e16..2343764677 100644 --- a/src/common/meta/src/kv_backend/memory.rs +++ b/src/common/meta/src/kv_backend/memory.rs @@ -16,20 +16,32 @@ use std::any::Any; use std::collections::btree_map::Entry; use std::collections::BTreeMap; use std::fmt::{Display, Formatter}; +use std::marker::PhantomData; +use std::ops::Range; use std::sync::RwLock; -use async_stream::stream; use async_trait::async_trait; +use common_error::prelude::ErrorExt; +use common_telemetry::timer; use serde::Serializer; -use crate::error::Error; -use crate::kv_backend::{Kv, KvBackend, ValueIter}; +use crate::kv_backend::txn::{Txn, TxnOp, TxnOpResponse, TxnRequest, TxnResponse}; +use crate::kv_backend::{KvBackend, TxnService}; +use crate::metrics::METRIC_META_TXN_REQUEST; +use crate::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, + BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, + DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, + RangeRequest, RangeResponse, +}; +use crate::rpc::KeyValue; -pub struct MemoryKvBackend { +pub struct MemoryKvBackend { kvs: RwLock, Vec>>, + _phantom: PhantomData, } -impl Display for MemoryKvBackend { +impl Display for MemoryKvBackend { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { let kvs = self.kvs.read().unwrap(); for (k, v) in kvs.iter() { @@ -42,93 +54,239 @@ impl Display for MemoryKvBackend { } } -impl Default for MemoryKvBackend { +impl Default for MemoryKvBackend { fn default() -> Self { Self { kvs: RwLock::new(BTreeMap::new()), + _phantom: PhantomData, } } } +impl MemoryKvBackend { + pub fn new() -> Self { + Self::default() + } + + pub fn clear(&self) { + let mut kvs = self.kvs.write().unwrap(); + kvs.clear(); + } +} + #[async_trait] -impl KvBackend for MemoryKvBackend { - type Error = Error; +impl KvBackend for MemoryKvBackend { + fn name(&self) -> &str { + "Memory" + } + + async fn range(&self, req: RangeRequest) -> Result { + let RangeRequest { + key, + range_end, + limit, + keys_only, + } = req; - fn range<'a, 'b>(&'a self, prefix: &[u8]) -> ValueIter<'b, Error> - where - 'a: 'b, - { let kvs = self.kvs.read().unwrap(); - let kvs = kvs.clone(); - let prefix = prefix.to_vec(); - Box::pin(stream!({ - for (k, v) in kvs.range(prefix.clone()..) { - if !k.starts_with(&prefix) { - break; + let iter: Box, &Vec)>> = if range_end.is_empty() { + Box::new(kvs.get_key_value(&key).into_iter()) + } else { + Box::new(kvs.range(key..range_end)) + }; + let mut kvs = iter + .map(|(k, v)| { + let key = k.clone(); + let value = if keys_only { vec![] } else { v.clone() }; + KeyValue { key, value } + }) + .collect::>(); + + let more = if limit > 0 && kvs.len() > limit as usize { + kvs.truncate(limit as usize); + true + } else { + false + }; + + Ok(RangeResponse { kvs, more }) + } + + async fn put(&self, req: PutRequest) -> Result { + let PutRequest { + key, + value, + prev_kv, + } = req; + + let mut kvs = self.kvs.write().unwrap(); + + let prev_kv = if prev_kv { + kvs.insert(key.clone(), value) + .map(|value| KeyValue { key, value }) + } else { + kvs.insert(key, value); + None + }; + + Ok(PutResponse { prev_kv }) + } + + async fn batch_put(&self, req: BatchPutRequest) -> Result { + let mut kvs = self.kvs.write().unwrap(); + + let mut prev_kvs = if req.prev_kv { + Vec::with_capacity(req.kvs.len()) + } else { + vec![] + }; + + for kv in req.kvs { + if req.prev_kv { + if let Some(value) = kvs.insert(kv.key.clone(), kv.value) { + prev_kvs.push(KeyValue { key: kv.key, value }); } - yield Ok(Kv(k.clone(), v.clone())); + } else { + kvs.insert(kv.key, kv.value); } - })) + } + + Ok(BatchPutResponse { prev_kvs }) } - async fn set(&self, key: &[u8], val: &[u8]) -> Result<(), Error> { - let mut kvs = self.kvs.write().unwrap(); - let _ = kvs.insert(key.to_vec(), val.to_vec()); - Ok(()) - } - - async fn compare_and_set( + async fn compare_and_put( &self, - key: &[u8], - expect: &[u8], - val: &[u8], - ) -> Result>>, Error> { - let key = key.to_vec(); - let val = val.to_vec(); + req: CompareAndPutRequest, + ) -> Result { + let CompareAndPutRequest { key, expect, value } = req; let mut kvs = self.kvs.write().unwrap(); + let existed = kvs.entry(key); - Ok(match existed { + let (success, prev_kv) = match existed { Entry::Vacant(e) => { - if expect.is_empty() { - let _ = e.insert(val); - Ok(()) - } else { - Err(None) + let expected = expect.is_empty(); + if expected { + let _ = e.insert(value); } + (expected, None) } Entry::Occupied(mut existed) => { - if existed.get() == expect { - let _ = existed.insert(val); - Ok(()) + let expected = existed.get() == &expect; + let prev_kv = if expected { + let _ = existed.insert(value); + None } else { - Err(Some(existed.get().clone())) - } + Some(KeyValue { + key: existed.key().clone(), + value: existed.get().clone(), + }) + }; + (expected, prev_kv) } + }; + + Ok(CompareAndPutResponse { success, prev_kv }) + } + + async fn delete_range( + &self, + req: DeleteRangeRequest, + ) -> Result { + let DeleteRangeRequest { + key, + range_end, + prev_kv, + } = req; + + let mut kvs = self.kvs.write().unwrap(); + + let prev_kvs = if range_end.is_empty() { + kvs.remove(&key) + .into_iter() + .map(|value| KeyValue { + key: key.clone(), + value, + }) + .collect::>() + } else { + let range = Range { + start: key, + end: range_end, + }; + kvs.drain_filter(|key, _| range.contains(key)) + .map(Into::into) + .collect::>() + }; + + Ok(DeleteRangeResponse { + deleted: prev_kvs.len() as i64, + prev_kvs: if prev_kv { prev_kvs } else { vec![] }, }) } - async fn delete_range(&self, key: &[u8], end: &[u8]) -> Result<(), Error> { + async fn batch_delete( + &self, + req: BatchDeleteRequest, + ) -> Result { let mut kvs = self.kvs.write().unwrap(); - if end.is_empty() { - let _ = kvs.remove(key); - } else { - let start = key.to_vec(); - let end = end.to_vec(); - let range = start..end; - kvs.retain(|k, _| !range.contains(k)); + let mut prev_kvs = if req.prev_kv { + Vec::with_capacity(req.keys.len()) + } else { + vec![] + }; + + for key in req.keys { + if req.prev_kv { + if let Some(value) = kvs.remove(&key) { + prev_kvs.push(KeyValue { key, value }); + } + } else { + kvs.remove(&key); + } } - Ok(()) + + Ok(BatchDeleteResponse { prev_kvs }) } - async fn move_value(&self, from_key: &[u8], to_key: &[u8]) -> Result<(), Error> { + async fn batch_get(&self, req: BatchGetRequest) -> Result { + let kvs = self.kvs.read().unwrap(); + + let kvs = req + .keys + .into_iter() + .filter_map(|key| { + kvs.get_key_value(&key).map(|(k, v)| KeyValue { + key: k.clone(), + value: v.clone(), + }) + }) + .collect::>(); + + Ok(BatchGetResponse { kvs }) + } + + async fn move_value(&self, req: MoveValueRequest) -> Result { + let MoveValueRequest { from_key, to_key } = req; + let mut kvs = self.kvs.write().unwrap(); - if let Some(v) = kvs.remove(from_key) { - let _ = kvs.insert(to_key.to_vec(), v); - } - Ok(()) + + let kv = if let Some(v) = kvs.remove(&from_key) { + kvs.insert(to_key, v.clone()); + Some(KeyValue { + key: from_key, + value: v, + }) + } else { + kvs.get(&to_key).map(|v| KeyValue { + key: to_key, + value: v.clone(), + }) + }; + + Ok(MoveValueResponse(kv)) } fn as_any(&self) -> &dyn Any { @@ -136,62 +294,383 @@ impl KvBackend for MemoryKvBackend { } } -#[cfg(test)] -mod tests { - use futures::TryStreamExt; +#[async_trait] +impl TxnService for MemoryKvBackend { + type Error = T; - use super::*; - - #[tokio::test] - async fn test_memory_kv_backend() { - let backend = MemoryKvBackend::default(); - - for i in 1..10 { - let key = format!("key{}", i); - let val = format!("val{}", i); - assert!(backend.set(key.as_bytes(), val.as_bytes()).await.is_ok()); - } - - let result = backend - .compare_and_set(b"hello", b"what", b"world") - .await - .unwrap(); - assert!(result.unwrap_err().is_none()); - - let result = backend - .compare_and_set(b"hello", b"", b"world") - .await - .unwrap(); - assert!(result.is_ok()); - - let result = backend - .compare_and_set(b"hello", b"world", b"greptime") - .await - .unwrap(); - assert!(result.is_ok()); - - let result = backend - .compare_and_set(b"hello", b"world", b"what") - .await - .unwrap(); - assert_eq!(result.unwrap_err().unwrap(), b"greptime"); - - assert!(backend.delete_range(b"key1", &[]).await.is_ok()); - assert!(backend.delete_range(b"key3", b"key9").await.is_ok()); - - assert!(backend.move_value(b"key9", b"key10").await.is_ok()); - - assert_eq!( - backend.to_string(), - r#"hello -> greptime -key10 -> val9 -key2 -> val2 -"# + async fn txn(&self, txn: Txn) -> Result { + let _timer = timer!( + METRIC_META_TXN_REQUEST, + &[("target", "memory"), ("op", "txn")] ); - let range = backend.range(b"key").try_collect::>().await.unwrap(); - assert_eq!(range.len(), 2); - assert_eq!(range[0], Kv(b"key10".to_vec(), b"val9".to_vec())); - assert_eq!(range[1], Kv(b"key2".to_vec(), b"val2".to_vec())); + let TxnRequest { + compare, + success, + failure, + } = txn.into(); + + let mut kvs = self.kvs.write().unwrap(); + + let succeeded = compare + .iter() + .all(|x| x.compare_with_value(kvs.get(&x.key))); + + let do_txn = |txn_op| match txn_op { + TxnOp::Put(key, value) => { + let prev_value = kvs.insert(key.clone(), value); + let prev_kv = prev_value.map(|value| KeyValue { key, value }); + TxnOpResponse::ResponsePut(PutResponse { prev_kv }) + } + + TxnOp::Get(key) => { + let value = kvs.get(&key); + let kvs = value + .into_iter() + .map(|value| KeyValue { + key: key.clone(), + value: value.clone(), + }) + .collect(); + TxnOpResponse::ResponseGet(RangeResponse { kvs, more: false }) + } + + TxnOp::Delete(key) => { + let prev_value = kvs.remove(&key); + let deleted = prev_value.as_ref().map(|x| x.len()).unwrap_or(0) as i64; + + let prev_kvs = prev_value + .into_iter() + .map(|value| KeyValue { + key: key.clone(), + value, + }) + .collect(); + TxnOpResponse::ResponseDelete(DeleteRangeResponse { deleted, prev_kvs }) + } + }; + + let responses: Vec<_> = if succeeded { success } else { failure } + .into_iter() + .map(do_txn) + .collect(); + + Ok(TxnResponse { + succeeded, + responses, + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicU8, Ordering}; + use std::sync::Arc; + + use super::*; + use crate::error::Error; + use crate::kv_backend::KvBackend; + use crate::rpc::store::{BatchGetRequest, BatchPutRequest}; + use crate::rpc::KeyValue; + use crate::util; + + async fn mock_mem_store_with_data() -> MemoryKvBackend { + let kv_store = MemoryKvBackend::::new(); + let kvs = mock_kvs(); + + assert!(kv_store + .batch_put(BatchPutRequest { + kvs, + ..Default::default() + }) + .await + .is_ok()); + + assert!(kv_store + .put(PutRequest { + key: b"key11".to_vec(), + value: b"val11".to_vec(), + ..Default::default() + }) + .await + .is_ok()); + + kv_store + } + + fn mock_kvs() -> Vec { + vec![ + KeyValue { + key: b"key1".to_vec(), + value: b"val1".to_vec(), + }, + KeyValue { + key: b"key2".to_vec(), + value: b"val2".to_vec(), + }, + KeyValue { + key: b"key3".to_vec(), + value: b"val3".to_vec(), + }, + ] + } + + #[tokio::test] + async fn test_put() { + let kv_store = mock_mem_store_with_data().await; + + let resp = kv_store + .put(PutRequest { + key: b"key11".to_vec(), + value: b"val12".to_vec(), + prev_kv: false, + }) + .await + .unwrap(); + assert!(resp.prev_kv.is_none()); + + let resp = kv_store + .put(PutRequest { + key: b"key11".to_vec(), + value: b"val13".to_vec(), + prev_kv: true, + }) + .await + .unwrap(); + let prev_kv = resp.prev_kv.unwrap(); + assert_eq!(b"key11", prev_kv.key()); + assert_eq!(b"val12", prev_kv.value()); + } + + #[tokio::test] + async fn test_range() { + let kv_store = mock_mem_store_with_data().await; + + let key = b"key1".to_vec(); + let range_end = util::get_prefix_end_key(b"key1"); + + let resp = kv_store + .range(RangeRequest { + key: key.clone(), + range_end: range_end.clone(), + limit: 0, + keys_only: false, + }) + .await + .unwrap(); + + assert_eq!(2, resp.kvs.len()); + assert_eq!(b"key1", resp.kvs[0].key()); + assert_eq!(b"val1", resp.kvs[0].value()); + assert_eq!(b"key11", resp.kvs[1].key()); + assert_eq!(b"val11", resp.kvs[1].value()); + + let resp = kv_store + .range(RangeRequest { + key: key.clone(), + range_end: range_end.clone(), + limit: 0, + keys_only: true, + }) + .await + .unwrap(); + + assert_eq!(2, resp.kvs.len()); + assert_eq!(b"key1", resp.kvs[0].key()); + assert_eq!(b"", resp.kvs[0].value()); + assert_eq!(b"key11", resp.kvs[1].key()); + assert_eq!(b"", resp.kvs[1].value()); + + let resp = kv_store + .range(RangeRequest { + key: key.clone(), + limit: 0, + keys_only: false, + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(1, resp.kvs.len()); + assert_eq!(b"key1", resp.kvs[0].key()); + assert_eq!(b"val1", resp.kvs[0].value()); + + let resp = kv_store + .range(RangeRequest { + key, + range_end, + limit: 1, + keys_only: false, + }) + .await + .unwrap(); + + assert_eq!(1, resp.kvs.len()); + assert_eq!(b"key1", resp.kvs[0].key()); + assert_eq!(b"val1", resp.kvs[0].value()); + } + + #[tokio::test] + async fn test_batch_get() { + let kv_store = mock_mem_store_with_data().await; + + let keys = vec![]; + let resp = kv_store.batch_get(BatchGetRequest { keys }).await.unwrap(); + + assert!(resp.kvs.is_empty()); + + let keys = vec![b"key10".to_vec()]; + let resp = kv_store.batch_get(BatchGetRequest { keys }).await.unwrap(); + + assert!(resp.kvs.is_empty()); + + let keys = vec![b"key1".to_vec(), b"key3".to_vec(), b"key4".to_vec()]; + let resp = kv_store.batch_get(BatchGetRequest { keys }).await.unwrap(); + + assert_eq!(2, resp.kvs.len()); + assert_eq!(b"key1", resp.kvs[0].key()); + assert_eq!(b"val1", resp.kvs[0].value()); + assert_eq!(b"key3", resp.kvs[1].key()); + assert_eq!(b"val3", resp.kvs[1].value()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_compare_and_put() { + let kv_store = Arc::new(MemoryKvBackend::::new()); + let success = Arc::new(AtomicU8::new(0)); + + let mut joins = vec![]; + for _ in 0..20 { + let kv_store_clone = kv_store.clone(); + let success_clone = success.clone(); + let join = tokio::spawn(async move { + let req = CompareAndPutRequest { + key: b"key".to_vec(), + expect: vec![], + value: b"val_new".to_vec(), + }; + let resp = kv_store_clone.compare_and_put(req).await.unwrap(); + if resp.success { + success_clone.fetch_add(1, Ordering::SeqCst); + } + }); + joins.push(join); + } + + for join in joins { + join.await.unwrap(); + } + + assert_eq!(1, success.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_delete_range() { + let kv_store = mock_mem_store_with_data().await; + + let req = DeleteRangeRequest { + key: b"key3".to_vec(), + range_end: vec![], + prev_kv: true, + }; + + let resp = kv_store.delete_range(req).await.unwrap(); + assert_eq!(1, resp.prev_kvs.len()); + assert_eq!(b"key3", resp.prev_kvs[0].key()); + assert_eq!(b"val3", resp.prev_kvs[0].value()); + + let resp = kv_store.get(b"key3").await.unwrap(); + assert!(resp.is_none()); + + let req = DeleteRangeRequest { + key: b"key2".to_vec(), + range_end: vec![], + prev_kv: false, + }; + + let resp = kv_store.delete_range(req).await.unwrap(); + assert!(resp.prev_kvs.is_empty()); + + let resp = kv_store.get(b"key2").await.unwrap(); + assert!(resp.is_none()); + + let key = b"key1".to_vec(); + let range_end = util::get_prefix_end_key(b"key1"); + + let req = DeleteRangeRequest { + key: key.clone(), + range_end: range_end.clone(), + prev_kv: true, + }; + let resp = kv_store.delete_range(req).await.unwrap(); + assert_eq!(2, resp.prev_kvs.len()); + + let req = RangeRequest { + key, + range_end, + ..Default::default() + }; + let resp = kv_store.range(req).await.unwrap(); + assert!(resp.kvs.is_empty()); + } + + #[tokio::test] + async fn test_move_value() { + let kv_store = mock_mem_store_with_data().await; + + let req = MoveValueRequest { + from_key: b"key1".to_vec(), + to_key: b"key111".to_vec(), + }; + + let resp = kv_store.move_value(req).await.unwrap(); + assert_eq!(b"key1", resp.0.as_ref().unwrap().key()); + assert_eq!(b"val1", resp.0.as_ref().unwrap().value()); + + let kv_store = mock_mem_store_with_data().await; + + let req = MoveValueRequest { + from_key: b"notexistkey".to_vec(), + to_key: b"key222".to_vec(), + }; + + let resp = kv_store.move_value(req).await.unwrap(); + assert!(resp.0.is_none()); + } + + #[tokio::test] + async fn test_batch_delete() { + let kv_store = mock_mem_store_with_data().await; + + assert!(kv_store.get(b"key1").await.unwrap().is_some()); + assert!(kv_store.get(b"key100").await.unwrap().is_none()); + + let req = BatchDeleteRequest { + keys: vec![b"key1".to_vec(), b"key100".to_vec()], + prev_kv: true, + }; + let resp = kv_store.batch_delete(req).await.unwrap(); + assert_eq!(1, resp.prev_kvs.len()); + assert_eq!( + vec![KeyValue { + key: b"key1".to_vec(), + value: b"val1".to_vec() + }], + resp.prev_kvs + ); + assert!(kv_store.get(b"key1").await.unwrap().is_none()); + + assert!(kv_store.get(b"key2").await.unwrap().is_some()); + assert!(kv_store.get(b"key3").await.unwrap().is_some()); + + let req = BatchDeleteRequest { + keys: vec![b"key2".to_vec(), b"key3".to_vec()], + prev_kv: false, + }; + let resp = kv_store.batch_delete(req).await.unwrap(); + assert!(resp.prev_kvs.is_empty()); + + assert!(kv_store.get(b"key2").await.unwrap().is_none()); + assert!(kv_store.get(b"key3").await.unwrap().is_none()); } } diff --git a/src/meta-srv/src/service/store/txn.rs b/src/common/meta/src/kv_backend/txn.rs similarity index 94% rename from src/meta-srv/src/service/store/txn.rs rename to src/common/meta/src/kv_backend/txn.rs index 4132b5d0a3..fe4a55e705 100644 --- a/src/meta-srv/src/service/store/txn.rs +++ b/src/common/meta/src/kv_backend/txn.rs @@ -12,15 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{DeleteRangeResponse, PutResponse, RangeResponse}; - -use crate::error::Result; - mod etcd; +use common_error::prelude::ErrorExt; + +use crate::rpc::store::{DeleteRangeResponse, PutResponse, RangeResponse}; + #[async_trait::async_trait] pub trait TxnService: Sync + Send { - async fn txn(&self, _txn: Txn) -> Result { + type Error: ErrorExt; + + async fn txn(&self, _txn: Txn) -> Result { unimplemented!("txn is not implemented") } } @@ -169,11 +171,14 @@ impl From for TxnRequest { #[cfg(test)] mod tests { - use api::v1::meta::{KeyValue, PutRequest}; + use std::sync::Arc; use super::*; - use crate::service::store::ext::KvStoreExt; - use crate::service::store::kv::KvStoreRef; + use crate::error::Error; + use crate::kv_backend::memory::MemoryKvBackend; + use crate::kv_backend::KvBackendRef; + use crate::rpc::store::PutRequest; + use crate::rpc::KeyValue; #[test] fn test_compare() { @@ -301,7 +306,7 @@ mod tests { async fn test_txn_compare_equal() { let kv_store = create_kv_store().await; let key = vec![101u8]; - let _ = kv_store.delete(key.clone(), false).await.unwrap(); + kv_store.delete(&key, false).await.unwrap(); let txn = Txn::new() .when(vec![Compare::with_not_exist_value( @@ -332,7 +337,7 @@ mod tests { async fn test_txn_compare_greater() { let kv_store = create_kv_store().await; let key = vec![102u8]; - let _ = kv_store.delete(key.clone(), false).await.unwrap(); + kv_store.delete(&key, false).await.unwrap(); let txn = Txn::new() .when(vec![Compare::with_not_exist_value( @@ -361,10 +366,9 @@ mod tests { assert_eq!( res, TxnOpResponse::ResponseGet(RangeResponse { - header: None, kvs: vec![KeyValue { key, - value: vec![1], + value: vec![1] }], more: false, }) @@ -375,7 +379,7 @@ mod tests { async fn test_txn_compare_less() { let kv_store = create_kv_store().await; let key = vec![103u8]; - let _ = kv_store.delete(vec![3], false).await.unwrap(); + kv_store.delete(&[3], false).await.unwrap(); let txn = Txn::new() .when(vec![Compare::with_not_exist_value( @@ -404,10 +408,9 @@ mod tests { assert_eq!( res, TxnOpResponse::ResponseGet(RangeResponse { - header: None, kvs: vec![KeyValue { key, - value: vec![2], + value: vec![2] }], more: false, }) @@ -418,7 +421,7 @@ mod tests { async fn test_txn_compare_not_equal() { let kv_store = create_kv_store().await; let key = vec![104u8]; - let _ = kv_store.delete(key.clone(), false).await.unwrap(); + kv_store.delete(&key, false).await.unwrap(); let txn = Txn::new() .when(vec![Compare::with_not_exist_value( @@ -447,18 +450,17 @@ mod tests { assert_eq!( res, TxnOpResponse::ResponseGet(RangeResponse { - header: None, kvs: vec![KeyValue { key, - value: vec![1], + value: vec![1] }], more: false, }) ); } - async fn create_kv_store() -> KvStoreRef { - std::sync::Arc::new(crate::service::store::memory::MemStore::new()) + async fn create_kv_store() -> KvBackendRef { + Arc::new(MemoryKvBackend::::new()) // TODO(jiachun): Add a feature to test against etcd in github CI // // The same test can be run against etcd by uncommenting the following line diff --git a/src/meta-srv/src/service/store/txn/etcd.rs b/src/common/meta/src/kv_backend/txn/etcd.rs similarity index 85% rename from src/meta-srv/src/service/store/txn/etcd.rs rename to src/common/meta/src/kv_backend/txn/etcd.rs index e2fc3c0315..ccee8b6062 100644 --- a/src/meta-srv/src/service/store/txn/etcd.rs +++ b/src/common/meta/src/kv_backend/txn/etcd.rs @@ -12,15 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{DeleteRangeResponse, PutResponse, RangeResponse}; use etcd_client::{ Compare as EtcdCompare, CompareOp as EtcdCompareOp, Txn as EtcdTxn, TxnOp as EtcdTxnOp, TxnOpResponse as EtcdTxnOpResponse, TxnResponse as EtcdTxnResponse, }; +use super::{Compare, CompareOp, Txn, TxnOp, TxnOpResponse, TxnResponse}; use crate::error::{self, Result}; -use crate::service::store::etcd_util::KvPair; -use crate::service::store::txn::{Compare, CompareOp, Txn, TxnOp, TxnOpResponse, TxnResponse}; +use crate::rpc::store::{DeleteRangeResponse, PutResponse, RangeResponse}; impl From for EtcdTxn { fn from(txn: Txn) -> Self { @@ -88,31 +87,25 @@ impl TryFrom for TxnOpResponse { fn try_from(op_resp: EtcdTxnOpResponse) -> Result { match op_resp { EtcdTxnOpResponse::Put(res) => { - let prev_kv = res.prev_key().map(KvPair::from_etcd_kv); - let put_res = PutResponse { - prev_kv, - ..Default::default() - }; + let prev_kv = res.prev_key().cloned().map(Into::into); + let put_res = PutResponse { prev_kv }; Ok(TxnOpResponse::ResponsePut(put_res)) } EtcdTxnOpResponse::Get(res) => { - let kvs = res.kvs().iter().map(KvPair::from_etcd_kv).collect(); - let range_res = RangeResponse { - kvs, - ..Default::default() - }; + let kvs = res.kvs().iter().cloned().map(Into::into).collect(); + let range_res = RangeResponse { kvs, more: false }; Ok(TxnOpResponse::ResponseGet(range_res)) } EtcdTxnOpResponse::Delete(res) => { let prev_kvs = res .prev_kvs() .iter() - .map(KvPair::from_etcd_kv) + .cloned() + .map(Into::into) .collect::>(); let delete_res = DeleteRangeResponse { prev_kvs, deleted: res.deleted(), - ..Default::default() }; Ok(TxnOpResponse::ResponseDelete(delete_res)) } diff --git a/src/common/meta/src/lib.rs b/src/common/meta/src/lib.rs index bd099570a3..aa93de1df1 100644 --- a/src/common/meta/src/lib.rs +++ b/src/common/meta/src/lib.rs @@ -12,15 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![feature(btree_drain_filter)] + pub mod error; pub mod heartbeat; pub mod ident; pub mod instruction; pub mod key; pub mod kv_backend; +pub mod metrics; pub mod peer; pub mod rpc; pub mod table_name; +pub mod util; pub type ClusterId = u64; pub type DatanodeId = u64; diff --git a/src/common/meta/src/metrics.rs b/src/common/meta/src/metrics.rs new file mode 100644 index 0000000000..01e3e3859f --- /dev/null +++ b/src/common/meta/src/metrics.rs @@ -0,0 +1,15 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub const METRIC_META_TXN_REQUEST: &str = "meta.txn_request"; diff --git a/src/common/meta/src/rpc.rs b/src/common/meta/src/rpc.rs index e87528e734..5d07408e2c 100644 --- a/src/common/meta/src/rpc.rs +++ b/src/common/meta/src/rpc.rs @@ -18,17 +18,14 @@ pub mod router; pub mod store; pub mod util; +use std::fmt::{Display, Formatter}; + use api::v1::meta::{KeyValue as PbKeyValue, ResponseHeader as PbResponseHeader}; #[derive(Debug, Clone)] pub struct ResponseHeader(PbResponseHeader); impl ResponseHeader { - #[inline] - pub(crate) fn new(header: PbResponseHeader) -> Self { - Self(header) - } - #[inline] pub fn protocol_version(&self) -> u64 { self.0.protocol_version @@ -56,33 +53,83 @@ impl ResponseHeader { } } -#[derive(Debug, Clone)] -pub struct KeyValue(PbKeyValue); +#[derive(Debug, Clone, PartialEq)] +pub struct KeyValue { + pub key: Vec, + pub value: Vec, +} + +impl From for PbKeyValue { + fn from(kv: KeyValue) -> Self { + Self { + key: kv.key, + value: kv.value, + } + } +} + +impl From for KeyValue { + fn from(kv: etcd_client::KeyValue) -> Self { + Self { + key: kv.key().to_vec(), + value: kv.value().to_vec(), + } + } +} + +impl From for (Vec, Vec) { + fn from(kv: KeyValue) -> Self { + (kv.key, kv.value) + } +} + +impl From<(Vec, Vec)> for KeyValue { + fn from(kv: (Vec, Vec)) -> Self { + Self { + key: kv.0, + value: kv.1, + } + } +} + +impl Display for KeyValue { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "({}, {})", + String::from_utf8_lossy(&self.key), + String::from_utf8_lossy(&self.value) + ) + } +} impl KeyValue { #[inline] - pub(crate) fn new(kv: PbKeyValue) -> Self { - Self(kv) + pub fn new(kv: PbKeyValue) -> Self { + Self { + key: kv.key, + value: kv.value, + } } #[inline] pub fn key(&self) -> &[u8] { - &self.0.key + &self.key } #[inline] pub fn take_key(&mut self) -> Vec { - std::mem::take(&mut self.0.key) + std::mem::take(&mut self.key) } #[inline] pub fn value(&self) -> &[u8] { - &self.0.value + &self.value } #[inline] pub fn take_value(&mut self) -> Vec { - std::mem::take(&mut self.0.value) + std::mem::take(&mut self.value) } } @@ -103,7 +150,7 @@ mod tests { }), }; - let header = ResponseHeader::new(pb_header); + let header = ResponseHeader(pb_header); assert_eq!(101, header.protocol_version()); assert_eq!(1, header.cluster_id()); assert_eq!(100, header.error_code()); diff --git a/src/common/meta/src/rpc/store.rs b/src/common/meta/src/rpc/store.rs index f3e8496d98..e1bdfdac2f 100644 --- a/src/common/meta/src/rpc/store.rs +++ b/src/common/meta/src/rpc/store.rs @@ -12,21 +12,23 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt::{Display, Formatter}; + use api::v1::meta::{ BatchDeleteRequest as PbBatchDeleteRequest, BatchDeleteResponse as PbBatchDeleteResponse, BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse, BatchPutRequest as PbBatchPutRequest, BatchPutResponse as PbBatchPutResponse, CompareAndPutRequest as PbCompareAndPutRequest, CompareAndPutResponse as PbCompareAndPutResponse, DeleteRangeRequest as PbDeleteRangeRequest, - DeleteRangeResponse as PbDeleteRangeResponse, KeyValue as PbKeyValue, - MoveValueRequest as PbMoveValueRequest, MoveValueResponse as PbMoveValueResponse, - PutRequest as PbPutRequest, PutResponse as PbPutResponse, RangeRequest as PbRangeRequest, - RangeResponse as PbRangeResponse, + DeleteRangeResponse as PbDeleteRangeResponse, MoveValueRequest as PbMoveValueRequest, + MoveValueResponse as PbMoveValueResponse, PutRequest as PbPutRequest, + PutResponse as PbPutResponse, RangeRequest as PbRangeRequest, RangeResponse as PbRangeResponse, + ResponseHeader as PbResponseHeader, }; use crate::error; use crate::error::Result; -use crate::rpc::{util, KeyValue, ResponseHeader}; +use crate::rpc::{util, KeyValue}; #[derive(Debug, Clone, Default)] pub struct RangeRequest { @@ -59,6 +61,30 @@ impl From for PbRangeRequest { } } +impl From for RangeRequest { + fn from(value: PbRangeRequest) -> Self { + Self { + key: value.key, + range_end: value.range_end, + limit: value.limit, + keys_only: value.keys_only, + } + } +} + +impl Display for RangeRequest { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "RangeRequest{{key: '{}', range_end: '{}', limit: {}, keys_only: {}}}", + String::from_utf8_lossy(&self.key), + String::from_utf8_lossy(&self.range_end), + self.limit, + self.keys_only + ) + } +} + impl RangeRequest { #[inline] pub fn new() -> Self { @@ -119,8 +145,26 @@ impl RangeRequest { } } -#[derive(Debug, Clone)] -pub struct RangeResponse(PbRangeResponse); +#[derive(Debug, Clone, PartialEq)] +pub struct RangeResponse { + pub kvs: Vec, + pub more: bool, +} + +impl Display for RangeResponse { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "RangeResponse{{kvs: [{}], more: {}}}", + self.kvs + .iter() + .map(|kv| kv.to_string()) + .collect::>() + .join(", "), + self.more + ) + } +} impl TryFrom for RangeResponse { type Error = error::Error; @@ -128,29 +172,25 @@ impl TryFrom for RangeResponse { fn try_from(pb: PbRangeResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self { + kvs: pb.kvs.into_iter().map(KeyValue::new).collect(), + more: pb.more, + }) } } impl RangeResponse { - #[inline] - pub fn new(res: PbRangeResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbRangeResponse { + PbRangeResponse { + header: Some(header), + kvs: self.kvs.into_iter().map(Into::into).collect(), + more: self.more, + } } #[inline] pub fn take_kvs(&mut self) -> Vec { - self.0.kvs.drain(..).map(KeyValue::new).collect() - } - - #[inline] - pub fn more(&self) -> bool { - self.0.more + self.kvs.drain(..).collect() } } @@ -177,6 +217,16 @@ impl From for PbPutRequest { } } +impl From for PutRequest { + fn from(value: PbPutRequest) -> Self { + Self { + key: value.key, + value: value.value, + prev_kv: value.prev_kv, + } + } +} + impl PutRequest { #[inline] pub fn new() -> Self { @@ -211,8 +261,10 @@ impl PutRequest { } } -#[derive(Debug, Clone)] -pub struct PutResponse(PbPutResponse); +#[derive(Debug, Clone, PartialEq)] +pub struct PutResponse { + pub prev_kv: Option, +} impl TryFrom for PutResponse { type Error = error::Error; @@ -220,27 +272,22 @@ impl TryFrom for PutResponse { fn try_from(pb: PbPutResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self { + prev_kv: pb.prev_kv.map(KeyValue::new), + }) } } impl PutResponse { - #[inline] - pub fn new(res: PbPutResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) - } - - #[inline] - pub fn take_prev_kv(&mut self) -> Option { - self.0.prev_kv.take().map(KeyValue::new) + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbPutResponse { + PbPutResponse { + header: Some(header), + prev_kv: self.prev_kv.map(Into::into), + } } } +#[derive(Clone)] pub struct BatchGetRequest { pub keys: Vec>, } @@ -254,6 +301,12 @@ impl From for PbBatchGetRequest { } } +impl From for BatchGetRequest { + fn from(value: PbBatchGetRequest) -> Self { + Self { keys: value.keys } + } +} + impl Default for BatchGetRequest { fn default() -> Self { Self::new() @@ -274,7 +327,23 @@ impl BatchGetRequest { } #[derive(Debug, Clone)] -pub struct BatchGetResponse(PbBatchGetResponse); +pub struct BatchGetResponse { + pub kvs: Vec, +} + +impl Display for BatchGetResponse { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "[{}]", + self.kvs + .iter() + .map(|kv| kv.to_string()) + .collect::>() + .join(", "), + ) + } +} impl TryFrom for BatchGetResponse { type Error = error::Error; @@ -282,30 +351,24 @@ impl TryFrom for BatchGetResponse { fn try_from(pb: PbBatchGetResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self(pb)) + Ok(Self { + kvs: pb.kvs.into_iter().map(KeyValue::new).collect(), + }) } } impl BatchGetResponse { - #[inline] - pub fn new(res: PbBatchGetResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) - } - - #[inline] - pub fn take_kvs(&mut self) -> Vec { - self.0.kvs.drain(..).map(KeyValue::new).collect() + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbBatchGetResponse { + PbBatchGetResponse { + header: Some(header), + kvs: self.kvs.into_iter().map(Into::into).collect(), + } } } #[derive(Debug, Clone, Default)] pub struct BatchPutRequest { - pub kvs: Vec, + pub kvs: Vec, /// If prev_kv is set, gets the previous key-value pairs before changing it. /// The previous key-value pairs will be returned in the batch put response. pub prev_kv: bool, @@ -315,12 +378,21 @@ impl From for PbBatchPutRequest { fn from(req: BatchPutRequest) -> Self { Self { header: None, - kvs: req.kvs, + kvs: req.kvs.into_iter().map(Into::into).collect(), prev_kv: req.prev_kv, } } } +impl From for BatchPutRequest { + fn from(value: PbBatchPutRequest) -> Self { + Self { + kvs: value.kvs.into_iter().map(KeyValue::new).collect(), + prev_kv: value.prev_kv, + } + } +} + impl BatchPutRequest { #[inline] pub fn new() -> Self { @@ -332,7 +404,7 @@ impl BatchPutRequest { #[inline] pub fn add_kv(mut self, key: impl Into>, value: impl Into>) -> Self { - self.kvs.push(PbKeyValue { + self.kvs.push(KeyValue { key: key.into(), value: value.into(), }); @@ -349,7 +421,9 @@ impl BatchPutRequest { } #[derive(Debug, Clone)] -pub struct BatchPutResponse(PbBatchPutResponse); +pub struct BatchPutResponse { + pub prev_kvs: Vec, +} impl TryFrom for BatchPutResponse { type Error = error::Error; @@ -357,24 +431,23 @@ impl TryFrom for BatchPutResponse { fn try_from(pb: PbBatchPutResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self { + prev_kvs: pb.prev_kvs.into_iter().map(KeyValue::new).collect(), + }) } } impl BatchPutResponse { - #[inline] - pub fn new(res: PbBatchPutResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbBatchPutResponse { + PbBatchPutResponse { + header: Some(header), + prev_kvs: self.prev_kvs.into_iter().map(Into::into).collect(), + } } #[inline] pub fn take_prev_kvs(&mut self) -> Vec { - self.0.prev_kvs.drain(..).map(KeyValue::new).collect() + self.prev_kvs.drain(..).collect() } } @@ -396,6 +469,15 @@ impl From for PbBatchDeleteRequest { } } +impl From for BatchDeleteRequest { + fn from(value: PbBatchDeleteRequest) -> Self { + Self { + keys: value.keys, + prev_kv: value.prev_kv, + } + } +} + impl BatchDeleteRequest { #[inline] pub fn new() -> Self { @@ -421,7 +503,9 @@ impl BatchDeleteRequest { } #[derive(Debug, Clone)] -pub struct BatchDeleteResponse(PbBatchDeleteResponse); +pub struct BatchDeleteResponse { + pub prev_kvs: Vec, +} impl TryFrom for BatchDeleteResponse { type Error = error::Error; @@ -429,24 +513,18 @@ impl TryFrom for BatchDeleteResponse { fn try_from(pb: PbBatchDeleteResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self { + prev_kvs: pb.prev_kvs.into_iter().map(KeyValue::new).collect(), + }) } } impl BatchDeleteResponse { - #[inline] - pub fn new(res: PbBatchDeleteResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) - } - - #[inline] - pub fn take_prev_kvs(&mut self) -> Vec { - self.0.prev_kvs.drain(..).map(KeyValue::new).collect() + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbBatchDeleteResponse { + PbBatchDeleteResponse { + header: Some(header), + prev_kvs: self.prev_kvs.into_iter().map(Into::into).collect(), + } } } @@ -471,6 +549,16 @@ impl From for PbCompareAndPutRequest { } } +impl From for CompareAndPutRequest { + fn from(value: PbCompareAndPutRequest) -> Self { + Self { + key: value.key, + expect: value.expect, + value: value.value, + } + } +} + impl CompareAndPutRequest { #[inline] pub fn new() -> Self { @@ -504,8 +592,11 @@ impl CompareAndPutRequest { } } -#[derive(Debug, Clone)] -pub struct CompareAndPutResponse(PbCompareAndPutResponse); +#[derive(Debug, Clone, Default)] +pub struct CompareAndPutResponse { + pub success: bool, + pub prev_kv: Option, +} impl TryFrom for CompareAndPutResponse { type Error = error::Error; @@ -513,29 +604,30 @@ impl TryFrom for CompareAndPutResponse { fn try_from(pb: PbCompareAndPutResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self { + success: pb.success, + prev_kv: pb.prev_kv.map(KeyValue::new), + }) } } impl CompareAndPutResponse { - #[inline] - pub fn new(res: PbCompareAndPutResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbCompareAndPutResponse { + PbCompareAndPutResponse { + header: Some(header), + success: self.success, + prev_kv: self.prev_kv.map(Into::into), + } } #[inline] pub fn is_success(&self) -> bool { - self.0.success + self.success } #[inline] pub fn take_prev_kv(&mut self) -> Option { - self.0.prev_kv.take().map(KeyValue::new) + self.prev_kv.take() } } @@ -571,6 +663,16 @@ impl From for PbDeleteRangeRequest { } } +impl From for DeleteRangeRequest { + fn from(value: PbDeleteRangeRequest) -> Self { + Self { + key: value.key, + range_end: value.range_end, + prev_kv: value.prev_kv, + } + } +} + impl DeleteRangeRequest { #[inline] pub fn new() -> Self { @@ -624,8 +726,11 @@ impl DeleteRangeRequest { } } -#[derive(Debug, Clone)] -pub struct DeleteRangeResponse(PbDeleteRangeResponse); +#[derive(Debug, Clone, PartialEq)] +pub struct DeleteRangeResponse { + pub deleted: i64, + pub prev_kvs: Vec, +} impl TryFrom for DeleteRangeResponse { type Error = error::Error; @@ -633,29 +738,30 @@ impl TryFrom for DeleteRangeResponse { fn try_from(pb: PbDeleteRangeResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self { + deleted: pb.deleted, + prev_kvs: pb.prev_kvs.into_iter().map(KeyValue::new).collect(), + }) } } impl DeleteRangeResponse { - #[inline] - pub fn new(res: PbDeleteRangeResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbDeleteRangeResponse { + PbDeleteRangeResponse { + header: Some(header), + deleted: self.deleted, + prev_kvs: self.prev_kvs.into_iter().map(Into::into).collect(), + } } #[inline] pub fn deleted(&self) -> i64 { - self.0.deleted + self.deleted } #[inline] pub fn take_prev_kvs(&mut self) -> Vec { - self.0.prev_kvs.drain(..).map(KeyValue::new).collect() + self.prev_kvs.drain(..).collect() } } @@ -678,6 +784,15 @@ impl From for PbMoveValueRequest { } } +impl From for MoveValueRequest { + fn from(value: PbMoveValueRequest) -> Self { + Self { + from_key: value.from_key, + to_key: value.to_key, + } + } +} + impl MoveValueRequest { #[inline] pub fn new(from_key: impl Into>, to_key: impl Into>) -> Self { @@ -689,7 +804,7 @@ impl MoveValueRequest { } #[derive(Debug, Clone)] -pub struct MoveValueResponse(PbMoveValueResponse); +pub struct MoveValueResponse(pub Option); impl TryFrom for MoveValueResponse { type Error = error::Error; @@ -697,24 +812,21 @@ impl TryFrom for MoveValueResponse { fn try_from(pb: PbMoveValueResponse) -> Result { util::check_response_header(pb.header.as_ref())?; - Ok(Self::new(pb)) + Ok(Self(pb.kv.map(KeyValue::new))) } } impl MoveValueResponse { - #[inline] - pub fn new(res: PbMoveValueResponse) -> Self { - Self(res) - } - - #[inline] - pub fn take_header(&mut self) -> Option { - self.0.header.take().map(ResponseHeader::new) + pub fn to_proto_resp(self, header: PbResponseHeader) -> PbMoveValueResponse { + PbMoveValueResponse { + header: Some(header), + kv: self.0.map(Into::into), + } } #[inline] pub fn take_kv(&mut self) -> Option { - self.0.kv.take().map(KeyValue::new) + self.0.take() } } @@ -784,9 +896,8 @@ mod tests { more: true, }; - let mut res = RangeResponse::new(pb_res); - assert!(res.take_header().is_none()); - assert!(res.more()); + let mut res: RangeResponse = pb_res.try_into().unwrap(); + assert!(res.more); let mut kvs = res.take_kvs(); let kv0 = kvs.get_mut(0).unwrap(); assert_eq!(b"k1".to_vec(), kv0.key().to_vec()); @@ -826,9 +937,8 @@ mod tests { }), }; - let mut res = PutResponse::new(pb_res); - assert!(res.take_header().is_none()); - let mut kv = res.take_prev_kv().unwrap(); + let res: PutResponse = pb_res.try_into().unwrap(); + let mut kv = res.prev_kv.unwrap(); assert_eq!(b"k1".to_vec(), kv.key().to_vec()); assert_eq!(b"k1".to_vec(), kv.take_key()); assert_eq!(b"v1".to_vec(), kv.value().to_vec()); @@ -859,12 +969,8 @@ mod tests { value: b"test_value1".to_vec(), }], }; - let mut res = BatchGetResponse::new(pb_res); - - assert!(res.take_header().is_none()); - - let kvs = res.take_kvs(); - + let res: BatchGetResponse = pb_res.try_into().unwrap(); + let kvs = res.kvs; assert_eq!(b"test_key1".as_slice(), kvs[0].key()); assert_eq!(b"test_value1".as_slice(), kvs[0].value()); } @@ -898,8 +1004,7 @@ mod tests { }], }; - let mut res = BatchPutResponse::new(pb_res); - assert!(res.take_header().is_none()); + let mut res: BatchPutResponse = pb_res.try_into().unwrap(); let kvs = res.take_prev_kvs(); assert_eq!(b"k1".to_vec(), kvs[0].key().to_vec()); assert_eq!(b"v1".to_vec(), kvs[0].value().to_vec()); @@ -931,9 +1036,8 @@ mod tests { }], }; - let mut res = BatchDeleteResponse::new(pb_res); - assert!(res.take_header().is_none()); - let kvs = res.take_prev_kvs(); + let res: BatchDeleteResponse = pb_res.try_into().unwrap(); + let kvs = res.prev_kvs; assert_eq!(b"k1".to_vec(), kvs[0].key().to_vec()); assert_eq!(b"v1".to_vec(), kvs[0].value().to_vec()); } @@ -969,8 +1073,7 @@ mod tests { }), }; - let mut res = CompareAndPutResponse::new(pb_res); - assert!(res.take_header().is_none()); + let mut res: CompareAndPutResponse = pb_res.try_into().unwrap(); let mut kv = res.take_prev_kv().unwrap(); assert_eq!(b"k1".to_vec(), kv.key().to_vec()); assert_eq!(b"k1".to_vec(), kv.take_key()); @@ -1026,7 +1129,6 @@ mod tests { }; let mut res: DeleteRangeResponse = pb_res.try_into().unwrap(); - assert!(res.take_header().is_none()); assert_eq!(2, res.deleted()); let mut kvs = res.take_prev_kvs(); let kv0 = kvs.get_mut(0).unwrap(); @@ -1064,8 +1166,7 @@ mod tests { }), }; - let mut res = MoveValueResponse::new(pb_res); - assert!(res.take_header().is_none()); + let mut res: MoveValueResponse = pb_res.try_into().unwrap(); let mut kv = res.take_kv().unwrap(); assert_eq!(b"k1".to_vec(), kv.key().to_vec()); assert_eq!(b"k1".to_vec(), kv.take_key()); diff --git a/src/meta-srv/src/util.rs b/src/common/meta/src/util.rs similarity index 100% rename from src/meta-srv/src/util.rs rename to src/common/meta/src/util.rs diff --git a/src/frontend/src/catalog.rs b/src/frontend/src/catalog.rs index 5973f66211..5fc45f87c3 100644 --- a/src/frontend/src/catalog.rs +++ b/src/frontend/src/catalog.rs @@ -34,11 +34,11 @@ use catalog::{ use client::client_manager::DatanodeClients; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME}; use common_error::prelude::BoxedError; -use common_meta::kv_backend::{Kv, KvBackendRef}; +use common_meta::kv_backend::KvBackendRef; +use common_meta::rpc::store::RangeRequest; +use common_meta::rpc::KeyValue; use common_meta::table_name::TableName; use common_telemetry::warn; -use futures::StreamExt; -use futures_util::TryStreamExt; use partition::manager::PartitionRuleManagerRef; use snafu::prelude::*; use table::table::numbers::NumbersTable; @@ -252,10 +252,17 @@ impl CatalogManager for FrontendCatalogManager { async fn catalog_names(&self) -> CatalogResult> { let key = build_catalog_prefix(); - let mut iter = self.backend.range(key.as_bytes()); + let req = RangeRequest::new().with_prefix(key.as_bytes()); + + let kvs = self + .backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; + let mut res = HashSet::new(); - while let Some(r) = iter.next().await { - let Kv(k, _) = r.context(TableMetadataManagerSnafu)?; + for KeyValue { key: k, value: _ } in kvs { let catalog_key = String::from_utf8_lossy(&k); if let Ok(key) = CatalogKey::parse(catalog_key.as_ref()) { let _ = res.insert(key.catalog_name); @@ -268,10 +275,17 @@ impl CatalogManager for FrontendCatalogManager { async fn schema_names(&self, catalog: &str) -> CatalogResult> { let key = build_schema_prefix(catalog); - let mut iter = self.backend.range(key.as_bytes()); + let req = RangeRequest::new().with_prefix(key.as_bytes()); + + let kvs = self + .backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs; + let mut res = HashSet::new(); - while let Some(r) = iter.next().await { - let Kv(k, _) = r.context(TableMetadataManagerSnafu)?; + for KeyValue { key: k, value: _ } in kvs { let key = SchemaKey::parse(String::from_utf8_lossy(&k)).context(InvalidCatalogValueSnafu)?; let _ = res.insert(key.schema_name); @@ -285,16 +299,23 @@ impl CatalogManager for FrontendCatalogManager { tables.push("numbers".to_string()); } let key = build_table_global_prefix(catalog, schema); - let iter = self.backend.range(key.as_bytes()); + let req = RangeRequest::new().with_prefix(key.as_bytes()); + + let iter = self + .backend + .range(req) + .await + .context(TableMetadataManagerSnafu)? + .kvs + .into_iter(); + let result = iter - .map(|r| { - let Kv(k, _) = r.context(TableMetadataManagerSnafu)?; + .map(|KeyValue { key: k, value: _ }| { let key = TableGlobalKey::parse(String::from_utf8_lossy(&k)) .context(InvalidCatalogValueSnafu)?; Ok(key.table_name) }) - .try_collect::>() - .await?; + .collect::>>()?; tables.extend(result); Ok(tables) } @@ -377,7 +398,7 @@ impl CatalogManager for FrontendCatalogManager { let Some(kv) = self.backend().get(table_global_key.to_string().as_bytes()).await.context(TableMetadataManagerSnafu)? else { return Ok(None); }; - let v = TableGlobalValue::from_bytes(kv.1).context(InvalidCatalogValueSnafu)?; + let v = TableGlobalValue::from_bytes(kv.value).context(InvalidCatalogValueSnafu)?; let table_info = Arc::new( v.table_info .try_into() diff --git a/src/frontend/src/instance/distributed/inserter.rs b/src/frontend/src/instance/distributed/inserter.rs index 5d00f6e17c..e59887b511 100644 --- a/src/frontend/src/instance/distributed/inserter.rs +++ b/src/frontend/src/instance/distributed/inserter.rs @@ -183,6 +183,7 @@ mod tests { use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackend, KvBackendRef}; + use common_meta::rpc::store::PutRequest; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -199,26 +200,20 @@ mod tests { catalog_name: DEFAULT_CATALOG_NAME.to_string(), } .to_string(); - backend - .set( - default_catalog.as_bytes(), - CatalogValue.as_bytes().unwrap().as_slice(), - ) - .await - .unwrap(); + let req = PutRequest::new() + .with_key(default_catalog.as_bytes()) + .with_value(CatalogValue.as_bytes().unwrap()); + backend.put(req).await.unwrap(); let default_schema = SchemaKey { catalog_name: DEFAULT_CATALOG_NAME.to_string(), schema_name: DEFAULT_SCHEMA_NAME.to_string(), } .to_string(); - backend - .set( - default_schema.as_bytes(), - SchemaValue.as_bytes().unwrap().as_slice(), - ) - .await - .unwrap(); + let req = PutRequest::new() + .with_key(default_schema.as_bytes()) + .with_value(SchemaValue.as_bytes().unwrap()); + backend.put(req).await.unwrap(); backend } @@ -257,13 +252,10 @@ mod tests { table_info: table_info.into(), }; - backend - .set( - table_global_key.to_string().as_bytes(), - table_global_value.as_bytes().unwrap().as_slice(), - ) - .await - .unwrap(); + let req = PutRequest::new() + .with_key(table_global_key.to_string().as_bytes()) + .with_value(table_global_value.as_bytes().unwrap()); + backend.put(req).await.unwrap(); } #[tokio::test] diff --git a/src/frontend/src/table.rs b/src/frontend/src/table.rs index 0377400f0d..84865adf0b 100644 --- a/src/frontend/src/table.rs +++ b/src/frontend/src/table.rs @@ -23,6 +23,7 @@ use catalog::helper::{TableGlobalKey, TableGlobalValue}; use client::Database; use common_error::prelude::BoxedError; use common_meta::key::TableRouteKey; +use common_meta::rpc::store::{MoveValueRequest, PutRequest}; use common_meta::table_name::TableName; use common_query::error::Result as QueryResult; use common_query::logical_plan::Expr; @@ -260,7 +261,7 @@ impl DistTable { .await .context(TableMetadataManagerSnafu)?; Ok(if let Some(raw) = raw { - Some(TableGlobalValue::from_bytes(raw.1).context(error::CatalogEntrySerdeSnafu)?) + Some(TableGlobalValue::from_bytes(raw.value).context(error::CatalogEntrySerdeSnafu)?) } else { None }) @@ -272,19 +273,26 @@ impl DistTable { value: TableGlobalValue, ) -> Result<()> { let value = value.as_bytes().context(error::CatalogEntrySerdeSnafu)?; - self.catalog_manager + let req = PutRequest::new() + .with_key(key.to_string().as_bytes()) + .with_value(value); + let _ = self + .catalog_manager .backend() - .set(key.to_string().as_bytes(), &value) + .put(req) .await - .context(TableMetadataManagerSnafu) + .context(TableMetadataManagerSnafu)?; + Ok(()) } async fn delete_table_global_value(&self, key: TableGlobalKey) -> Result<()> { - self.catalog_manager + let _ = self + .catalog_manager .backend() - .delete(key.to_string().as_bytes()) + .delete(key.to_string().as_bytes(), false) .await - .context(TableMetadataManagerSnafu) + .context(TableMetadataManagerSnafu)?; + Ok(()) } async fn move_table_route_value( @@ -311,9 +319,10 @@ impl DistTable { } .to_string(); + let req = MoveValueRequest::new(old_key.as_bytes(), new_key.as_bytes()); self.catalog_manager .backend() - .move_value(old_key.as_bytes(), new_key.as_bytes()) + .move_value(req) .await .context(TableMetadataManagerSnafu)?; diff --git a/src/meta-client/Cargo.toml b/src/meta-client/Cargo.toml index 0945cadd37..43cf2b222d 100644 --- a/src/meta-client/Cargo.toml +++ b/src/meta-client/Cargo.toml @@ -12,7 +12,7 @@ common-error = { path = "../common/error" } common-grpc = { path = "../common/grpc" } common-telemetry = { path = "../common/telemetry" } common-meta = { path = "../common/meta" } -etcd-client = "0.11" +etcd-client.workspace = true rand.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index efd8fe6ec5..5b871a99b4 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -739,7 +739,7 @@ mod tests { .with_key(tc.key("key")) .with_value(b"value".to_vec()); let res = tc.client.put(req).await; - assert!(res.unwrap().take_prev_kv().is_none()); + assert!(res.unwrap().prev_kv.is_none()); } #[tokio::test] @@ -752,14 +752,14 @@ mod tests { .with_value(b"value".to_vec()) .with_prev_kv(); let res = tc.client.put(req).await; - assert!(res.unwrap().take_prev_kv().is_none()); + assert!(res.unwrap().prev_kv.is_none()); let req = PutRequest::new() .with_key(key.as_slice()) .with_value(b"value1".to_vec()) .with_prev_kv(); let res = tc.client.put(req).await; - let mut kv = res.unwrap().take_prev_kv().unwrap(); + let mut kv = res.unwrap().prev_kv.unwrap(); assert_eq!(key, kv.take_key()); assert_eq!(b"value".to_vec(), kv.take_value()); } @@ -794,16 +794,14 @@ mod tests { for i in 0..256 { req = req.add_key(tc.key(&format!("key-{}", i))); } - let mut res = tc.client.batch_get(req).await.unwrap(); - - assert_eq!(10, res.take_kvs().len()); + let res = tc.client.batch_get(req).await.unwrap(); + assert_eq!(10, res.kvs.len()); let req = BatchGetRequest::default() .add_key(tc.key("key-1")) .add_key(tc.key("key-999")); - let mut res = tc.client.batch_get(req).await.unwrap(); - - assert_eq!(1, res.take_kvs().len()); + let res = tc.client.batch_get(req).await.unwrap(); + assert_eq!(1, res.kvs.len()); } #[tokio::test] @@ -867,7 +865,9 @@ mod tests { let res = tc.client.compare_and_put(req).await; let mut res = res.unwrap(); assert!(res.is_success()); - assert_eq!(b"value".to_vec(), res.take_prev_kv().unwrap().take_value()); + + // If compare-and-put is success, previous value doesn't need to be returned. + assert!(res.take_prev_kv().is_none()); } #[tokio::test] diff --git a/src/meta-srv/Cargo.toml b/src/meta-srv/Cargo.toml index 27f8d89103..dd935c9664 100644 --- a/src/meta-srv/Cargo.toml +++ b/src/meta-srv/Cargo.toml @@ -25,7 +25,7 @@ common-telemetry = { path = "../common/telemetry" } common-time = { path = "../common/time" } dashmap = "5.4" derive_builder = "0.12" -etcd-client = "0.11" +etcd-client.workspace = true futures.workspace = true h2 = "0.3" http-body = "0.4" diff --git a/src/meta-srv/examples/kv_store.rs b/src/meta-srv/examples/kv_store.rs index bae7a1a996..aa5dbbc19a 100644 --- a/src/meta-srv/examples/kv_store.rs +++ b/src/meta-srv/examples/kv_store.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{DeleteRangeRequest, PutRequest, RangeRequest}; +use common_meta::rpc::store::{DeleteRangeRequest, PutRequest, RangeRequest}; use meta_srv::service::store::etcd::EtcdStore; use tracing::{event, subscriber, Level}; use tracing_subscriber::FmtSubscriber; @@ -31,7 +31,6 @@ async fn run() { key: b"key1".to_vec(), value: b"value1".to_vec(), prev_kv: true, - ..Default::default() }; let res = kv_store.put(put_req).await; event!(Level::INFO, "put result: {:#?}", res); diff --git a/src/meta-srv/src/cluster.rs b/src/meta-srv/src/cluster.rs index fa25a21eab..1cc7d80cdf 100644 --- a/src/meta-srv/src/cluster.rs +++ b/src/meta-srv/src/cluster.rs @@ -18,18 +18,22 @@ use std::time::Duration; use api::v1::meta::cluster_client::ClusterClient; use api::v1::meta::{ - BatchGetRequest, BatchGetResponse, KeyValue, RangeRequest, RangeResponse, ResponseHeader, + BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse, + RangeRequest as PbRangeRequest, RangeResponse as PbRangeResponse, ResponseHeader, }; use common_grpc::channel_manager::ChannelManager; +use common_meta::rpc::store::{BatchGetRequest, RangeRequest}; +use common_meta::rpc::KeyValue; +use common_meta::util; use common_telemetry::warn; use derive_builder::Builder; use snafu::{ensure, OptionExt, ResultExt}; +use crate::error; use crate::error::{match_for_io_error, Result}; use crate::keys::{StatKey, StatValue, DN_STAT_PREFIX}; use crate::metasrv::ElectionRef; use crate::service::store::kv::ResettableKvStoreRef; -use crate::{error, util}; pub type MetaPeerClientRef = Arc; @@ -112,13 +116,13 @@ impl MetaPeerClient { .get(&leader_addr) .context(error::CreateChannelSnafu)?; - let request = tonic::Request::new(RangeRequest { + let request = tonic::Request::new(PbRangeRequest { key, range_end, ..Default::default() }); - let response: RangeResponse = ClusterClient::new(channel) + let response: PbRangeResponse = ClusterClient::new(channel) .range(request) .await .context(error::RangeSnafu)? @@ -126,16 +130,13 @@ impl MetaPeerClient { check_resp_header(&response.header, Context { addr: &leader_addr })?; - Ok(response.kvs) + Ok(response.kvs.into_iter().map(KeyValue::new).collect()) } // Get kv information from the leader's in_mem kv store pub async fn batch_get(&self, keys: Vec>) -> Result> { if self.is_leader() { - let request = BatchGetRequest { - keys, - ..Default::default() - }; + let request = BatchGetRequest { keys }; return self.in_memory.batch_get(request).await.map(|resp| resp.kvs); } @@ -175,12 +176,12 @@ impl MetaPeerClient { .get(&leader_addr) .context(error::CreateChannelSnafu)?; - let request = tonic::Request::new(BatchGetRequest { + let request = tonic::Request::new(PbBatchGetRequest { keys, ..Default::default() }); - let response: BatchGetResponse = ClusterClient::new(channel) + let response: PbBatchGetResponse = ClusterClient::new(channel) .batch_get(request) .await .context(error::BatchGetSnafu)? @@ -188,7 +189,7 @@ impl MetaPeerClient { check_resp_header(&response.header, Context { addr: &leader_addr })?; - Ok(response.kvs) + Ok(response.kvs.into_iter().map(KeyValue::new).collect()) } // Check if the meta node is a leader node. @@ -240,7 +241,8 @@ fn need_retry(error: &error::Error) -> bool { #[cfg(test)] mod tests { - use api::v1::meta::{Error, ErrorCode, KeyValue, ResponseHeader}; + use api::v1::meta::{Error, ErrorCode, ResponseHeader}; + use common_meta::rpc::KeyValue; use super::{check_resp_header, to_stat_kv_map, Context}; use crate::error; diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 8bfb8b071e..98c3b52c9e 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -409,8 +409,11 @@ pub enum Error { source: common_meta::error::Error, }, - #[snafu(display("Etcd txn got an error: {err_msg}"))] - EtcdTxnOpResponse { err_msg: String, location: Location }, + #[snafu(display("Failed to convert Etcd txn object: {source}"))] + ConvertEtcdTxnObject { + source: common_meta::error::Error, + location: Location, + }, // this error is used for custom error mapping // please do not delete it @@ -490,7 +493,6 @@ impl ErrorExt for Error { | Error::InvalidTxnResult { .. } | Error::InvalidUtf8Value { .. } | Error::UnexpectedInstructionReply { .. } - | Error::EtcdTxnOpResponse { .. } | Error::Unexpected { .. } | Error::Txn { .. } | Error::TableIdChanged { .. } => StatusCode::Unexpected, @@ -507,9 +509,11 @@ impl ErrorExt for Error { Error::RegionFailoverCandidatesNotFound { .. } => StatusCode::RuntimeResourcesExhausted, Error::RegisterProcedureLoader { source, .. } => source.status_code(), - Error::TableRouteConversion { source, .. } | Error::ConvertProtoData { source, .. } => { - source.status_code() - } + + Error::TableRouteConversion { source, .. } + | Error::ConvertProtoData { source, .. } + | Error::ConvertEtcdTxnObject { source, .. } => source.status_code(), + Error::Other { source, .. } => source.status_code(), } } diff --git a/src/meta-srv/src/handler/keep_lease_handler.rs b/src/meta-srv/src/handler/keep_lease_handler.rs index 3d1a1ad307..814d7f4828 100644 --- a/src/meta-srv/src/handler/keep_lease_handler.rs +++ b/src/meta-srv/src/handler/keep_lease_handler.rs @@ -12,7 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{HeartbeatRequest, PutRequest, Role}; +use api::v1::meta::{HeartbeatRequest, Role}; +use common_meta::rpc::store::PutRequest; use common_telemetry::{trace, warn}; use common_time::util as time_util; diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index 8ab335d9bf..28893f0397 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -14,7 +14,8 @@ use std::cmp::Ordering; -use api::v1::meta::{HeartbeatRequest, PutRequest, Role}; +use api::v1::meta::{HeartbeatRequest, Role}; +use common_meta::rpc::store::PutRequest; use common_telemetry::warn; use dashmap::DashMap; @@ -147,7 +148,6 @@ mod tests { use crate::keys::StatKey; use crate::sequence::Sequence; use crate::service::store::cached_kv::LeaderCachedKvStore; - use crate::service::store::ext::KvStoreExt; use crate::service::store::memory::MemStore; #[tokio::test] @@ -184,7 +184,8 @@ mod tests { cluster_id: 3, node_id: 101, }; - let res = ctx.in_memory.get(key.try_into().unwrap()).await.unwrap(); + let key: Vec = key.try_into().unwrap(); + let res = ctx.in_memory.get(&key).await.unwrap(); let kv = res.unwrap(); let key: StatKey = kv.key.clone().try_into().unwrap(); assert_eq!(3, key.cluster_id); @@ -195,7 +196,9 @@ mod tests { assert_eq!(Some(1), val.stats[0].region_num); handle_request_many_times(ctx.clone(), &handler, 10).await; - let res = ctx.in_memory.get(key.try_into().unwrap()).await.unwrap(); + + let key: Vec = key.try_into().unwrap(); + let res = ctx.in_memory.get(&key).await.unwrap(); let kv = res.unwrap(); let val: StatValue = kv.value.try_into().unwrap(); // refresh every 10 stats diff --git a/src/meta-srv/src/lease.rs b/src/meta-srv/src/lease.rs index 377bb7ffb9..c62d894ec1 100644 --- a/src/meta-srv/src/lease.rs +++ b/src/meta-srv/src/lease.rs @@ -14,12 +14,12 @@ use std::collections::HashMap; +use common_meta::util; use common_time::util as time_util; use crate::cluster::MetaPeerClientRef; use crate::error::Result; use crate::keys::{LeaseKey, LeaseValue, DN_LEASE_PREFIX}; -use crate::util; pub async fn alive_datanodes( cluster_id: u64, diff --git a/src/meta-srv/src/lib.rs b/src/meta-srv/src/lib.rs index af1314197e..9d9e0c28f6 100644 --- a/src/meta-srv/src/lib.rs +++ b/src/meta-srv/src/lib.rs @@ -40,6 +40,4 @@ pub mod table_routes; #[cfg(test)] mod test_util; -pub mod util; - pub use crate::error::Result; diff --git a/src/meta-srv/src/metadata_service.rs b/src/meta-srv/src/metadata_service.rs index db61d40d74..f27020a5ba 100644 --- a/src/meta-srv/src/metadata_service.rs +++ b/src/meta-srv/src/metadata_service.rs @@ -14,9 +14,9 @@ use std::sync::Arc; -use api::v1::meta::CompareAndPutRequest; use async_trait::async_trait; use catalog::helper::{CatalogKey, CatalogValue, SchemaKey, SchemaValue}; +use common_meta::rpc::store::CompareAndPutRequest; use common_telemetry::{info, timer}; use metrics::increment_counter; use snafu::{ensure, ResultExt}; @@ -80,7 +80,6 @@ impl MetadataService for DefaultMetadataService { value: CatalogValue {} .as_bytes() .context(error::InvalidCatalogValueSnafu)?, - ..Default::default() }; let resp = kv_store.compare_and_put(req).await?; @@ -96,7 +95,6 @@ impl MetadataService for DefaultMetadataService { value: SchemaValue {} .as_bytes() .context(error::InvalidCatalogValueSnafu)?, - ..Default::default() }; let resp = kv_store.compare_and_put(req).await?; @@ -124,7 +122,6 @@ mod tests { use catalog::helper::{CatalogKey, SchemaKey}; use super::{DefaultMetadataService, MetadataService}; - use crate::service::store::ext::KvStoreExt; use crate::service::store::kv::KvStoreRef; use crate::service::store::memory::MemStore; @@ -156,11 +153,9 @@ mod tests { .to_string() .into(); - let result = kv_store.get(key.clone()).await.unwrap(); - + let result = kv_store.get(&key).await.unwrap(); let kv = result.unwrap(); - - assert_eq!(key, kv.key); + assert_eq!(key, kv.key()); let key: Vec = SchemaKey { catalog_name: "catalog".to_string(), @@ -169,10 +164,8 @@ mod tests { .to_string() .into(); - let result = kv_store.get(key.clone()).await.unwrap(); - + let result = kv_store.get(&key).await.unwrap(); let kv = result.unwrap(); - - assert_eq!(key, kv.key); + assert_eq!(key, kv.key()); } } diff --git a/src/meta-srv/src/metrics.rs b/src/meta-srv/src/metrics.rs index 13cdd98c02..cac6598991 100644 --- a/src/meta-srv/src/metrics.rs +++ b/src/meta-srv/src/metrics.rs @@ -15,7 +15,6 @@ pub(crate) const METRIC_META_CREATE_CATALOG: &str = "meta.create_catalog"; pub(crate) const METRIC_META_CREATE_SCHEMA: &str = "meta.create_schema"; pub(crate) const METRIC_META_KV_REQUEST: &str = "meta.kv_request"; -pub(crate) const METRIC_META_TXN_REQUEST: &str = "meta.txn_request"; pub(crate) const METRIC_META_ROUTE_REQUEST: &str = "meta.route_request"; pub(crate) const METRIC_META_HEARTBEAT_CONNECTION_NUM: &str = "meta.heartbeat_connection_num"; pub(crate) const METRIC_META_HANDLER_EXECUTE: &str = "meta.handler_execute"; diff --git a/src/meta-srv/src/procedure/create_table.rs b/src/meta-srv/src/procedure/create_table.rs index 52e062dcb2..f3a98eb225 100644 --- a/src/meta-srv/src/procedure/create_table.rs +++ b/src/meta-srv/src/procedure/create_table.rs @@ -19,6 +19,7 @@ use client::Database; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_meta::key::TableRouteKey; +use common_meta::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp}; use common_meta::rpc::ddl::CreateTableTask; use common_meta::rpc::router::TableRoute; use common_meta::table_name::TableName; @@ -34,7 +35,6 @@ use super::utils::{handle_request_datanode_error, handle_retry_error}; use crate::ddl::DdlContext; use crate::error::{self, Result}; use crate::service::router::create_table_global_value; -use crate::service::store::txn::{Compare, CompareOp, Txn, TxnOp}; use crate::table_routes::get_table_global_value; // TODO(weny): removes in following PRs. diff --git a/src/meta-srv/src/procedure/drop_table.rs b/src/meta-srv/src/procedure/drop_table.rs index f36a83183e..76b2dffc3e 100644 --- a/src/meta-srv/src/procedure/drop_table.rs +++ b/src/meta-srv/src/procedure/drop_table.rs @@ -21,6 +21,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_meta::ident::TableIdent; use common_meta::instruction::Instruction; +use common_meta::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp}; use common_meta::rpc::ddl::DropTableTask; use common_meta::rpc::router::TableRoute; use common_meta::table_name::TableName; @@ -40,7 +41,6 @@ use crate::error; use crate::error::Result; use crate::procedure::utils::{build_table_route_value, handle_request_datanode_error}; use crate::service::mailbox::BroadcastChannel; -use crate::service::store::txn::{Compare, CompareOp, Txn, TxnOp}; use crate::table_routes::fetch_table; pub struct DropTableProcedure { context: DdlContext, diff --git a/src/meta-srv/src/procedure/region_failover.rs b/src/meta-srv/src/procedure/region_failover.rs index f291748515..167dfe170b 100644 --- a/src/meta-srv/src/procedure/region_failover.rs +++ b/src/meta-srv/src/procedure/region_failover.rs @@ -45,7 +45,6 @@ use crate::error::{Error, RegisterProcedureLoaderSnafu, Result}; use crate::lock::DistLockRef; use crate::metasrv::{SelectorContext, SelectorRef}; use crate::service::mailbox::MailboxRef; -use crate::service::store::ext::KvStoreExt; const OPEN_REGION_MESSAGE_TIMEOUT: Duration = Duration::from_secs(30); const CLOSE_REGION_MESSAGE_TIMEOUT: Duration = Duration::from_secs(2); @@ -208,7 +207,7 @@ impl RegionFailoverManager { let table_global_value = self .selector_ctx .kv_store - .get(table_global_key.to_raw_key()) + .get(&table_global_key.to_raw_key()) .await?; Ok(table_global_value.is_some()) } diff --git a/src/meta-srv/src/procedure/state_store.rs b/src/meta-srv/src/procedure/state_store.rs index 5745f10c00..bbb4baf99c 100644 --- a/src/meta-srv/src/procedure/state_store.rs +++ b/src/meta-srv/src/procedure/state_store.rs @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{BatchDeleteRequest, PutRequest, RangeRequest}; use async_stream::try_stream; use async_trait::async_trait; use common_error::prelude::BoxedError; +use common_meta::rpc::store::{BatchDeleteRequest, PutRequest, RangeRequest}; +use common_meta::util; use common_procedure::error::{ CorruptedDataSnafu, DeleteStatesSnafu, ListStateSnafu, PutStateSnafu, }; @@ -24,7 +25,6 @@ use common_procedure::Result; use snafu::ResultExt; use crate::service::store::kv::KvStoreRef; -use crate::util; const PROCEDURE_PREFIX: &str = "/__procedure__/"; diff --git a/src/meta-srv/src/sequence.rs b/src/meta-srv/src/sequence.rs index 747e488fb8..8dbd0771d2 100644 --- a/src/meta-srv/src/sequence.rs +++ b/src/meta-srv/src/sequence.rs @@ -15,7 +15,7 @@ use std::ops::Range; use std::sync::Arc; -use api::v1::meta::CompareAndPutRequest; +use common_meta::rpc::store::CompareAndPutRequest; use snafu::{ensure, OptionExt}; use tokio::sync::Mutex; @@ -120,7 +120,6 @@ impl Inner { key: key.to_vec(), expect, value: value.to_vec(), - ..Default::default() }; let res = self.generator.compare_and_put(req).await?; @@ -156,16 +155,20 @@ impl Inner { #[cfg(test)] mod tests { + use std::any::Any; use std::sync::Arc; - use api::v1::meta::{ + use common_meta::kv_backend::{KvBackend, TxnService}; + use common_meta::rpc::store::{ BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, CompareAndPutResponse, DeleteRangeRequest, + DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, + RangeRequest, RangeResponse, }; use super::*; - use crate::service::store::kv::KvStore; + use crate::error::Error; use crate::service::store::memory::MemStore; - use crate::service::store::txn::TxnService; #[tokio::test] async fn test_sequence() { @@ -200,28 +203,25 @@ mod tests { async fn test_sequence_force_quit() { struct Noop; - impl TxnService for Noop {} + impl TxnService for Noop { + type Error = Error; + } #[async_trait::async_trait] - impl KvStore for Noop { - async fn range( - &self, - _: api::v1::meta::RangeRequest, - ) -> Result { + impl KvBackend for Noop { + fn name(&self) -> &str { + "Noop" + } + + async fn range(&self, _: RangeRequest) -> Result { unreachable!() } - async fn put( - &self, - _: api::v1::meta::PutRequest, - ) -> Result { + async fn put(&self, _: PutRequest) -> Result { unreachable!() } - async fn batch_put( - &self, - _: api::v1::meta::BatchPutRequest, - ) -> Result { + async fn batch_put(&self, _: BatchPutRequest) -> Result { unreachable!() } @@ -232,27 +232,25 @@ mod tests { async fn compare_and_put( &self, _: CompareAndPutRequest, - ) -> Result { - Ok(api::v1::meta::CompareAndPutResponse::default()) + ) -> Result { + Ok(CompareAndPutResponse::default()) } - async fn delete_range( - &self, - _: api::v1::meta::DeleteRangeRequest, - ) -> Result { + async fn delete_range(&self, _: DeleteRangeRequest) -> Result { unreachable!() } - async fn move_value( - &self, - _: api::v1::meta::MoveValueRequest, - ) -> Result { + async fn move_value(&self, _: MoveValueRequest) -> Result { unreachable!() } async fn batch_delete(&self, _: BatchDeleteRequest) -> Result { unreachable!() } + + fn as_any(&self) -> &dyn Any { + self + } } let kv_store = Arc::new(Noop {}); diff --git a/src/meta-srv/src/service/admin/meta.rs b/src/meta-srv/src/service/admin/meta.rs index 4e9923d222..6bcb6c97d5 100644 --- a/src/meta-srv/src/service/admin/meta.rs +++ b/src/meta-srv/src/service/admin/meta.rs @@ -14,18 +14,18 @@ use std::collections::HashMap; -use api::v1::meta::{RangeRequest, RangeResponse}; use catalog::helper::{ build_catalog_prefix, build_schema_prefix, build_table_global_prefix, TABLE_GLOBAL_KEY_PREFIX, }; +use common_meta::rpc::store::{RangeRequest, RangeResponse}; +use common_meta::util; use snafu::{OptionExt, ResultExt}; use tonic::codegen::http; +use crate::error; use crate::error::Result; use crate::service::admin::HttpHandler; -use crate::service::store::ext::KvStoreExt; use crate::service::store::kv::KvStoreRef; -use crate::{error, util}; pub struct CatalogsHandler { pub kv_store: KvStoreRef, @@ -104,7 +104,7 @@ impl HttpHandler for TableHandler { })?; let table_key = format!("{TABLE_GLOBAL_KEY_PREFIX}-{table_name}"); - let response = self.kv_store.get(table_key.into_bytes()).await?; + let response = self.kv_store.get(table_key.as_bytes()).await?; let mut value: String = "Not found result".to_string(); if let Some(key_value) = response { value = String::from_utf8(key_value.value).context(error::InvalidUtf8ValueSnafu)?; @@ -161,11 +161,11 @@ async fn get_keys_by_prefix(key_prefix: String, kv_store: &KvStoreRef) -> Result mod tests { use std::sync::Arc; - use api::v1::meta::PutRequest; use catalog::helper::{ build_catalog_prefix, build_schema_prefix, build_table_global_prefix, CatalogKey, SchemaKey, TableGlobalKey, }; + use common_meta::rpc::store::PutRequest; use crate::service::admin::meta::get_keys_by_prefix; use crate::service::store::kv::KvStoreRef; @@ -184,7 +184,7 @@ mod tests { .put(PutRequest { key: catalog.to_string().as_bytes().to_vec(), value: "".as_bytes().to_vec(), - ..Default::default() + prev_kv: false, }) .await .is_ok()); @@ -197,7 +197,7 @@ mod tests { .put(PutRequest { key: schema.to_string().as_bytes().to_vec(), value: "".as_bytes().to_vec(), - ..Default::default() + prev_kv: false, }) .await .is_ok()); @@ -216,7 +216,7 @@ mod tests { .put(PutRequest { key: table1.to_string().as_bytes().to_vec(), value: "".as_bytes().to_vec(), - ..Default::default() + prev_kv: false, }) .await .is_ok()); @@ -224,7 +224,7 @@ mod tests { .put(PutRequest { key: table2.to_string().as_bytes().to_vec(), value: "".as_bytes().to_vec(), - ..Default::default() + prev_kv: false, }) .await .is_ok()); diff --git a/src/meta-srv/src/service/admin/route.rs b/src/meta-srv/src/service/admin/route.rs index d1ad65ab29..0768670282 100644 --- a/src/meta-srv/src/service/admin/route.rs +++ b/src/meta-srv/src/service/admin/route.rs @@ -14,16 +14,18 @@ use std::collections::HashMap; -use api::v1::meta::{RangeRequest, RangeResponse, TableRouteValue}; +use api::v1::meta::TableRouteValue; use common_meta::key::TABLE_ROUTE_PREFIX; +use common_meta::rpc::store::{RangeRequest, RangeResponse}; +use common_meta::util; use prost::Message; use snafu::{OptionExt, ResultExt}; use tonic::codegen::http; use super::HttpHandler; +use crate::error; use crate::error::Result; use crate::service::store::kv::KvStoreRef; -use crate::{error, util}; pub struct RouteHandler { pub kv_store: KvStoreRef, diff --git a/src/meta-srv/src/service/cluster.rs b/src/meta-srv/src/service/cluster.rs index 257cf36d15..ffe6d61fb0 100644 --- a/src/meta-srv/src/service/cluster.rs +++ b/src/meta-srv/src/service/cluster.rs @@ -13,8 +13,8 @@ // limitations under the License. use api::v1::meta::{ - cluster_server, BatchGetRequest, BatchGetResponse, Error, RangeRequest, RangeResponse, - ResponseHeader, + cluster_server, BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse, + Error, RangeRequest as PbRangeRequest, RangeResponse as PbRangeResponse, ResponseHeader, }; use common_telemetry::warn; use tonic::{Request, Response}; @@ -24,10 +24,10 @@ use crate::service::GrpcResult; #[async_trait::async_trait] impl cluster_server::Cluster for MetaSrv { - async fn batch_get(&self, req: Request) -> GrpcResult { + async fn batch_get(&self, req: Request) -> GrpcResult { if !self.is_leader() { let is_not_leader = ResponseHeader::failed(0, Error::is_not_leader()); - let resp = BatchGetResponse { + let resp = PbBatchGetResponse { header: Some(is_not_leader), ..Default::default() }; @@ -36,22 +36,17 @@ impl cluster_server::Cluster for MetaSrv { return Ok(Response::new(resp)); } - let kvs = self.in_memory().batch_get(req.into_inner()).await?.kvs; + let req = req.into_inner().into(); + let resp = self.in_memory().batch_get(req).await?; - let success = ResponseHeader::success(0); - - let get_resp = BatchGetResponse { - kvs, - header: Some(success), - }; - - Ok(Response::new(get_resp)) + let resp = resp.to_proto_resp(ResponseHeader::success(0)); + Ok(Response::new(resp)) } - async fn range(&self, req: Request) -> GrpcResult { + async fn range(&self, req: Request) -> GrpcResult { if !self.is_leader() { let is_not_leader = ResponseHeader::failed(0, Error::is_not_leader()); - let resp = RangeResponse { + let resp = PbRangeResponse { header: Some(is_not_leader), ..Default::default() }; @@ -60,10 +55,11 @@ impl cluster_server::Cluster for MetaSrv { return Ok(Response::new(resp)); } - let req = req.into_inner(); + let req = req.into_inner().into(); let res = self.in_memory().range(req).await?; - Ok(Response::new(res)) + let resp = res.to_proto_resp(ResponseHeader::success(0)); + Ok(Response::new(resp)) } } diff --git a/src/meta-srv/src/service/router.rs b/src/meta-srv/src/service/router.rs index 7837fc8e3b..faf7edd3bd 100644 --- a/src/meta-srv/src/service/router.rs +++ b/src/meta-srv/src/service/router.rs @@ -15,12 +15,13 @@ use std::collections::HashMap; use api::v1::meta::{ - router_server, BatchPutRequest, CreateRequest, DeleteRequest, Error, KeyValue, Peer, PeerDict, - Region, RegionRoute, ResponseHeader, RouteRequest, RouteResponse, Table, TableRoute, - TableRouteValue, + router_server, CreateRequest, DeleteRequest, Error, Peer, PeerDict, Region, RegionRoute, + ResponseHeader, RouteRequest, RouteResponse, Table, TableRoute, TableRouteValue, }; use catalog::helper::{TableGlobalKey, TableGlobalValue}; use common_meta::key::TableRouteKey; +use common_meta::rpc::store::BatchPutRequest; +use common_meta::rpc::KeyValue; use common_meta::table_name::TableName; use common_telemetry::{timer, warn}; use snafu::{ensure, OptionExt, ResultExt}; @@ -33,7 +34,6 @@ use crate::lock::{keys, DistLockGuard}; use crate::metasrv::{Context, MetaSrv, SelectorContext, SelectorRef}; use crate::metrics::METRIC_META_ROUTE_REQUEST; use crate::sequence::SequenceRef; -use crate::service::store::ext::KvStoreExt; use crate::service::GrpcResult; use crate::table_routes::{ fetch_tables, get_table_global_value, remove_table_global_value, remove_table_route_value, @@ -66,7 +66,7 @@ impl router_server::Router for MetaSrv { .to_string() .into_bytes(); ensure!( - self.kv_store().get(table_global_key).await?.is_none(), + self.kv_store().get(&table_global_key).await?.is_none(), error::TableAlreadyExistsSnafu { table_name: table_name.to_string(), } @@ -238,7 +238,6 @@ async fn handle_create( }, ], prev_kv: true, - ..Default::default() }; let resp = ctx.kv_store.batch_put(req).await?; diff --git a/src/meta-srv/src/service/store.rs b/src/meta-srv/src/service/store.rs index 8b184f10df..bafee9190b 100644 --- a/src/meta-srv/src/service/store.rs +++ b/src/meta-srv/src/service/store.rs @@ -15,85 +15,252 @@ pub mod cached_kv; pub mod etcd; pub(crate) mod etcd_util; -pub mod ext; pub mod kv; pub mod memory; -pub mod txn; use api::v1::meta::{ - store_server, BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, - BatchPutRequest, BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, - DeleteRangeRequest, DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, - PutResponse, RangeRequest, RangeResponse, + store_server, BatchDeleteRequest as PbBatchDeleteRequest, + BatchDeleteResponse as PbBatchDeleteResponse, BatchGetRequest as PbBatchGetRequest, + BatchGetResponse as PbBatchGetResponse, BatchPutRequest as PbBatchPutRequest, + BatchPutResponse as PbBatchPutResponse, CompareAndPutRequest as PbCompareAndPutRequest, + CompareAndPutResponse as PbCompareAndPutResponse, DeleteRangeRequest as PbDeleteRangeRequest, + DeleteRangeResponse as PbDeleteRangeResponse, MoveValueRequest as PbMoveValueRequest, + MoveValueResponse as PbMoveValueResponse, PutRequest as PbPutRequest, + PutResponse as PbPutResponse, RangeRequest as PbRangeRequest, RangeResponse as PbRangeResponse, + ResponseHeader, }; +use common_meta::rpc::store::{ + BatchDeleteRequest, BatchGetRequest, BatchPutRequest, CompareAndPutRequest, DeleteRangeRequest, + MoveValueRequest, PutRequest, RangeRequest, +}; +use common_telemetry::timer; +use snafu::OptionExt; use tonic::{Request, Response}; +use crate::error::MissingRequestHeaderSnafu; use crate::metasrv::MetaSrv; +use crate::metrics::METRIC_META_KV_REQUEST; use crate::service::GrpcResult; #[async_trait::async_trait] impl store_server::Store for MetaSrv { - async fn range(&self, req: Request) -> GrpcResult { + async fn range(&self, req: Request) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "range".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: RangeRequest = req.into(); + let res = self.kv_store().range(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } - async fn put(&self, req: Request) -> GrpcResult { + async fn put(&self, req: Request) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "put".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: PutRequest = req.into(); + let res = self.kv_store().put(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } - async fn batch_get(&self, req: Request) -> GrpcResult { + async fn batch_get(&self, req: Request) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "batch_get".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: BatchGetRequest = req.into(); + let res = self.kv_store().batch_get(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } - async fn batch_put(&self, req: Request) -> GrpcResult { + async fn batch_put(&self, req: Request) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "batch_pub".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: BatchPutRequest = req.into(); + let res = self.kv_store().batch_put(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } async fn batch_delete( &self, - req: Request, - ) -> GrpcResult { + req: Request, + ) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "batch_delete".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: BatchDeleteRequest = req.into(); + let res = self.kv_store().batch_delete(req).await?; + + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } async fn compare_and_put( &self, - req: Request, - ) -> GrpcResult { + req: Request, + ) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "compare_and_put".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: CompareAndPutRequest = req.into(); + let res = self.kv_store().compare_and_put(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } async fn delete_range( &self, - req: Request, - ) -> GrpcResult { + req: Request, + ) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "delete_range".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: DeleteRangeRequest = req.into(); + let res = self.kv_store().delete_range(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } - async fn move_value(&self, req: Request) -> GrpcResult { + async fn move_value( + &self, + req: Request, + ) -> GrpcResult { let req = req.into_inner(); + + let cluster_id = req + .header + .as_ref() + .context(MissingRequestHeaderSnafu)? + .cluster_id; + + let _timer = timer!( + METRIC_META_KV_REQUEST, + &[ + ("target", self.kv_store().name().to_string()), + ("op", "move_value".to_string()), + ("cluster_id", cluster_id.to_string()), + ] + ); + + let req: MoveValueRequest = req.into(); + let res = self.kv_store().move_value(req).await?; + let res = res.to_proto_resp(ResponseHeader::success(cluster_id)); Ok(Response::new(res)) } } @@ -122,7 +289,8 @@ mod tests { async fn test_range() { let meta_srv = new_meta_srv().await; - let req = RangeRequest::default(); + let mut req = RangeRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.range(req.into_request()).await; let _ = res.unwrap(); @@ -132,7 +300,8 @@ mod tests { async fn test_put() { let meta_srv = new_meta_srv().await; - let req = PutRequest::default(); + let mut req = PutRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.put(req.into_request()).await; let _ = res.unwrap(); @@ -142,7 +311,8 @@ mod tests { async fn test_batch_get() { let meta_srv = new_meta_srv().await; - let req = BatchGetRequest::default(); + let mut req = BatchGetRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.batch_get(req.into_request()).await; let _ = res.unwrap(); @@ -152,7 +322,8 @@ mod tests { async fn test_batch_put() { let meta_srv = new_meta_srv().await; - let req = BatchPutRequest::default(); + let mut req = BatchPutRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.batch_put(req.into_request()).await; let _ = res.unwrap(); @@ -162,7 +333,8 @@ mod tests { async fn test_batch_delete() { let meta_srv = new_meta_srv().await; - let req = BatchDeleteRequest::default(); + let mut req = BatchDeleteRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.batch_delete(req.into_request()).await; let _ = res.unwrap(); @@ -172,7 +344,8 @@ mod tests { async fn test_compare_and_put() { let meta_srv = new_meta_srv().await; - let req = CompareAndPutRequest::default(); + let mut req = CompareAndPutRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.compare_and_put(req.into_request()).await; let _ = res.unwrap(); @@ -182,7 +355,8 @@ mod tests { async fn test_delete_range() { let meta_srv = new_meta_srv().await; - let req = DeleteRangeRequest::default(); + let mut req = DeleteRangeRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.delete_range(req.into_request()).await; let _ = res.unwrap(); @@ -192,7 +366,8 @@ mod tests { async fn test_move_value() { let meta_srv = new_meta_srv().await; - let req = MoveValueRequest::default(); + let mut req = MoveValueRequest::default(); + req.set_header((1, 1), Role::Datanode); let res = meta_srv.move_value(req.into_request()).await; let _ = res.unwrap(); diff --git a/src/meta-srv/src/service/store/cached_kv.rs b/src/meta-srv/src/service/store/cached_kv.rs index 8637670a74..79ecf36992 100644 --- a/src/meta-srv/src/service/store/cached_kv.rs +++ b/src/meta-srv/src/service/store/cached_kv.rs @@ -12,22 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::any::Any; use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use api::v1::meta::{ +use common_meta::kv_backend::txn::{Txn, TxnOp, TxnRequest, TxnResponse}; +use common_meta::kv_backend::{KvBackend, TxnService}; +use common_meta::rpc::store::{ BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, - DeleteRangeResponse, KeyValue, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, + DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, RangeRequest, RangeResponse, }; +use common_meta::rpc::KeyValue; -use crate::error::Result; -use crate::service::store::ext::KvStoreExt; -use crate::service::store::kv::{KvStore, KvStoreRef, ResettableKvStore, ResettableKvStoreRef}; +use crate::error::{Error, Result}; +use crate::service::store::kv::{KvStoreRef, ResettableKvStore, ResettableKvStoreRef}; use crate::service::store::memory::MemStore; -use crate::service::store::txn::{Txn, TxnOp, TxnRequest, TxnResponse, TxnService}; pub type CheckLeaderRef = Arc; @@ -57,15 +59,18 @@ pub struct LeaderCachedKvStore { store: KvStoreRef, cache: ResettableKvStoreRef, version: AtomicUsize, + name: String, } impl LeaderCachedKvStore { pub fn new(check_leader: CheckLeaderRef, store: KvStoreRef) -> Self { + let name = format!("LeaderCached({})", store.name()); Self { check_leader, store, cache: Arc::new(MemStore::new()), version: AtomicUsize::new(0), + name, } } @@ -82,7 +87,7 @@ impl LeaderCachedKvStore { #[inline] async fn invalid_key(&self, key: Vec) -> Result<()> { - let _ = self.cache.delete(key, false).await?; + let _ = self.cache.delete(&key, false).await?; Ok(()) } @@ -110,7 +115,11 @@ impl LeaderCachedKvStore { } #[async_trait::async_trait] -impl KvStore for LeaderCachedKvStore { +impl KvBackend for LeaderCachedKvStore { + fn name(&self) -> &str { + &self.name + } + async fn range(&self, req: RangeRequest) -> Result { if !self.is_leader() { return self.store.range(req).await; @@ -186,16 +195,13 @@ impl KvStore for LeaderCachedKvStore { .filter(|key| !hit_keys.contains(*key)) .cloned() .collect::>(); - let remote_req = BatchGetRequest { - keys: missed_keys, - ..Default::default() - }; + let remote_req = BatchGetRequest { keys: missed_keys }; let ver = self.get_version(); let remote_res = self.store.batch_get(remote_req).await?; let put_req = BatchPutRequest { - kvs: remote_res.kvs.clone(), + kvs: remote_res.kvs.clone().into_iter().map(Into::into).collect(), ..Default::default() }; let _ = self.cache.batch_put(put_req).await?; @@ -204,7 +210,7 @@ impl KvStore for LeaderCachedKvStore { let keys = remote_res .kvs .iter() - .map(|kv| kv.key.clone()) + .map(|kv| kv.key().to_vec()) .collect::>(); self.invalid_keys(keys).await?; } @@ -291,10 +297,16 @@ impl KvStore for LeaderCachedKvStore { self.invalid_keys(vec![from_key, to_key]).await?; Ok(res) } + + fn as_any(&self) -> &dyn Any { + self + } } #[async_trait::async_trait] impl TxnService for LeaderCachedKvStore { + type Error = Error; + async fn txn(&self, txn: Txn) -> Result { if !self.is_leader() { return self.store.txn(txn).await; @@ -338,7 +350,7 @@ impl ResettableKvStore for LeaderCachedKvStore { #[cfg(test)] mod tests { - use api::v1::meta::KeyValue; + use common_meta::rpc::KeyValue; use super::*; use crate::service::store::memory::MemStore; @@ -364,23 +376,19 @@ mod tests { }; let _ = inner_store.put(put_req).await.unwrap(); - let cached_value = inner_cache.get(key.clone()).await.unwrap(); + let cached_value = inner_cache.get(&key).await.unwrap(); assert!(cached_value.is_none()); - let cached_value = cached_store.get(key.clone()).await.unwrap().unwrap(); - assert_eq!(cached_value.value, value); + let cached_value = cached_store.get(&key).await.unwrap().unwrap(); + assert_eq!(cached_value.value(), value); - let cached_value = inner_cache.get(key.clone()).await.unwrap().unwrap(); - assert_eq!(cached_value.value, value); + let cached_value = inner_cache.get(&key).await.unwrap().unwrap(); + assert_eq!(cached_value.value(), value); - let res = cached_store - .delete(key.clone(), true) - .await - .unwrap() - .unwrap(); - assert_eq!(res.value, value); + let res = cached_store.delete(&key, true).await.unwrap().unwrap(); + assert_eq!(res.value(), value); - let cached_value = inner_cache.get(key.clone()).await.unwrap(); + let cached_value = inner_cache.get(&key).await.unwrap(); assert!(cached_value.is_none()); } @@ -409,10 +417,7 @@ mod tests { .map(|i| format!("test_key_{}", i).into_bytes()) .collect::>(); - let batch_get_req = BatchGetRequest { - keys, - ..Default::default() - }; + let batch_get_req = BatchGetRequest { keys }; let cached_values = inner_cache.batch_get(batch_get_req.clone()).await.unwrap(); assert!(cached_values.kvs.is_empty()); @@ -451,10 +456,7 @@ mod tests { let keys = (1..5) .map(|i| format!("test_key_{}", i).into_bytes()) .collect::>(); - let batch_get_req = BatchGetRequest { - keys, - ..Default::default() - }; + let batch_get_req = BatchGetRequest { keys }; let cached_values = inner_cache.batch_get(batch_get_req.clone()).await.unwrap(); assert_eq!(cached_values.kvs.len(), 4); diff --git a/src/meta-srv/src/service/store/etcd.rs b/src/meta-srv/src/service/store/etcd.rs index f41dd7bc58..841a130887 100644 --- a/src/meta-srv/src/service/store/etcd.rs +++ b/src/meta-srv/src/service/store/etcd.rs @@ -12,15 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::any::Any; use std::sync::Arc; -use api::v1::meta::{ +use common_error::prelude::*; +use common_meta::kv_backend::txn::{Txn as KvTxn, TxnResponse as KvTxnResponse}; +use common_meta::kv_backend::{KvBackend, TxnService}; +use common_meta::metrics::METRIC_META_TXN_REQUEST; +use common_meta::rpc::store::{ BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, - DeleteRangeResponse, KeyValue, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, - RangeRequest, RangeResponse, ResponseHeader, + DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, + RangeRequest, RangeResponse, }; -use common_error::prelude::*; +use common_meta::rpc::KeyValue; use common_telemetry::{timer, warn}; use etcd_client::{ Client, Compare, CompareOp, DeleteOptions, GetOptions, PutOptions, Txn, TxnOp, TxnOpResponse, @@ -28,11 +33,9 @@ use etcd_client::{ }; use crate::error; -use crate::error::Result; -use crate::metrics::{METRIC_META_KV_REQUEST, METRIC_META_TXN_REQUEST}; +use crate::error::{ConvertEtcdTxnObjectSnafu, Error, Result}; use crate::service::store::etcd_util::KvPair; -use crate::service::store::kv::{KvStore, KvStoreRef}; -use crate::service::store::txn::TxnService; +use crate::service::store::kv::KvStoreRef; // Maximum number of operations permitted in a transaction. // The etcd default configuration's `--max-txn-ops` is 128. @@ -89,22 +92,13 @@ impl EtcdStore { } #[async_trait::async_trait] -impl KvStore for EtcdStore { - async fn range(&self, req: RangeRequest) -> Result { - let Get { - cluster_id, - key, - options, - } = req.try_into()?; +impl KvBackend for EtcdStore { + fn name(&self) -> &str { + "Etcd" + } - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "range".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); + async fn range(&self, req: RangeRequest) -> Result { + let Get { key, options } = req.try_into()?; let res = self .client @@ -119,9 +113,7 @@ impl KvStore for EtcdStore { .map(KvPair::from_etcd_kv) .collect::>(); - let header = Some(ResponseHeader::success(cluster_id)); Ok(RangeResponse { - header, kvs, more: res.more(), }) @@ -129,21 +121,11 @@ impl KvStore for EtcdStore { async fn put(&self, req: PutRequest) -> Result { let Put { - cluster_id, key, value, options, } = req.try_into()?; - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "put".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); - let res = self .client .kv_client() @@ -152,26 +134,11 @@ impl KvStore for EtcdStore { .context(error::EtcdFailedSnafu)?; let prev_kv = res.prev_key().map(KvPair::from_etcd_kv); - - let header = Some(ResponseHeader::success(cluster_id)); - Ok(PutResponse { header, prev_kv }) + Ok(PutResponse { prev_kv }) } async fn batch_get(&self, req: BatchGetRequest) -> Result { - let BatchGet { - cluster_id, - keys, - options, - } = req.try_into()?; - - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "batch_get".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); + let BatchGet { keys, options } = req.try_into()?; let get_ops: Vec<_> = keys .into_iter() @@ -192,25 +159,11 @@ impl KvStore for EtcdStore { } } - let header = Some(ResponseHeader::success(cluster_id)); - Ok(BatchGetResponse { header, kvs }) + Ok(BatchGetResponse { kvs }) } async fn batch_put(&self, req: BatchPutRequest) -> Result { - let BatchPut { - cluster_id, - kvs, - options, - } = req.try_into()?; - - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "batch_put".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); + let BatchPut { kvs, options } = req.try_into()?; let put_ops = kvs .into_iter() @@ -233,25 +186,11 @@ impl KvStore for EtcdStore { } } - let header = Some(ResponseHeader::success(cluster_id)); - Ok(BatchPutResponse { header, prev_kvs }) + Ok(BatchPutResponse { prev_kvs }) } async fn batch_delete(&self, req: BatchDeleteRequest) -> Result { - let BatchDelete { - cluster_id, - keys, - options, - } = req.try_into()?; - - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "batch_delete".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); + let BatchDelete { keys, options } = req.try_into()?; let mut prev_kvs = Vec::with_capacity(keys.len()); @@ -275,28 +214,17 @@ impl KvStore for EtcdStore { } } - let header = Some(ResponseHeader::success(cluster_id)); - Ok(BatchDeleteResponse { header, prev_kvs }) + Ok(BatchDeleteResponse { prev_kvs }) } async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result { let CompareAndPut { - cluster_id, key, expect, value, put_options, } = req.try_into()?; - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "compare_and_put".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); - let compare = if expect.is_empty() { // create if absent // revision 0 means key was not exist @@ -333,29 +261,11 @@ impl KvStore for EtcdStore { _ => unreachable!(), }; - let header = Some(ResponseHeader::success(cluster_id)); - Ok(CompareAndPutResponse { - header, - success, - prev_kv, - }) + Ok(CompareAndPutResponse { success, prev_kv }) } async fn delete_range(&self, req: DeleteRangeRequest) -> Result { - let Delete { - cluster_id, - key, - options, - } = req.try_into()?; - - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "delete_range".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); + let Delete { key, options } = req.try_into()?; let res = self .client @@ -370,9 +280,7 @@ impl KvStore for EtcdStore { .map(KvPair::from_etcd_kv) .collect::>(); - let header = Some(ResponseHeader::success(cluster_id)); Ok(DeleteRangeResponse { - header, deleted: res.deleted(), prev_kvs, }) @@ -380,24 +288,13 @@ impl KvStore for EtcdStore { async fn move_value(&self, req: MoveValueRequest) -> Result { let MoveValue { - cluster_id, from_key, to_key, delete_options, } = req.try_into()?; - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[ - ("target", "etcd".to_string()), - ("op", "move_value".to_string()), - ("cluster_id", cluster_id.to_string()) - ] - ); - let mut client = self.client.kv_client(); - let header = Some(ResponseHeader::success(cluster_id)); // TODO(jiachun): Maybe it's better to let the users control it in the request const MAX_RETRIES: usize = 8; for _ in 0..MAX_RETRIES { @@ -442,16 +339,14 @@ impl KvStore for EtcdStore { for op_res in txn_res.op_responses() { match op_res { TxnOpResponse::Get(res) => { - return Ok(MoveValueResponse { - header, - kv: res.kvs().first().map(KvPair::from_etcd_kv), - }); + return Ok(MoveValueResponse( + res.kvs().first().map(KvPair::from_etcd_kv), + )); } TxnOpResponse::Delete(res) => { - return Ok(MoveValueResponse { - header, - kv: res.prev_kvs().first().map(KvPair::from_etcd_kv), - }); + return Ok(MoveValueResponse( + res.prev_kvs().first().map(KvPair::from_etcd_kv), + )); } _ => {} } @@ -463,14 +358,17 @@ impl KvStore for EtcdStore { } .fail() } + + fn as_any(&self) -> &dyn Any { + self + } } #[async_trait::async_trait] impl TxnService for EtcdStore { - async fn txn( - &self, - txn: crate::service::store::txn::Txn, - ) -> Result { + type Error = Error; + + async fn txn(&self, txn: KvTxn) -> Result { let _timer = timer!( METRIC_META_TXN_REQUEST, &[("target", "etcd".to_string()), ("op", "txn".to_string()),] @@ -483,12 +381,11 @@ impl TxnService for EtcdStore { .txn(etcd_txn) .await .context(error::EtcdFailedSnafu)?; - txn_res.try_into() + txn_res.try_into().context(ConvertEtcdTxnObjectSnafu) } } struct Get { - cluster_id: u64, key: Vec, options: Option, } @@ -498,7 +395,6 @@ impl TryFrom for Get { fn try_from(req: RangeRequest) -> Result { let RangeRequest { - header, key, range_end, limit, @@ -519,7 +415,6 @@ impl TryFrom for Get { } Ok(Get { - cluster_id: header.map_or(0, |h| h.cluster_id), key, options: Some(options), }) @@ -527,7 +422,6 @@ impl TryFrom for Get { } struct Put { - cluster_id: u64, key: Vec, value: Vec, options: Option, @@ -538,7 +432,6 @@ impl TryFrom for Put { fn try_from(req: PutRequest) -> Result { let PutRequest { - header, key, value, prev_kv, @@ -550,7 +443,6 @@ impl TryFrom for Put { } Ok(Put { - cluster_id: header.map_or(0, |h| h.cluster_id), key, value, options: Some(options), @@ -559,7 +451,6 @@ impl TryFrom for Put { } struct BatchGet { - cluster_id: u64, keys: Vec>, options: Option, } @@ -568,12 +459,11 @@ impl TryFrom for BatchGet { type Error = error::Error; fn try_from(req: BatchGetRequest) -> Result { - let BatchGetRequest { header, keys } = req; + let BatchGetRequest { keys } = req; let options = GetOptions::default(); Ok(BatchGet { - cluster_id: header.map_or(0, |h| h.cluster_id), keys, options: Some(options), }) @@ -581,7 +471,6 @@ impl TryFrom for BatchGet { } struct BatchPut { - cluster_id: u64, kvs: Vec, options: Option, } @@ -590,11 +479,7 @@ impl TryFrom for BatchPut { type Error = error::Error; fn try_from(req: BatchPutRequest) -> Result { - let BatchPutRequest { - header, - kvs, - prev_kv, - } = req; + let BatchPutRequest { kvs, prev_kv } = req; let mut options = PutOptions::default(); if prev_kv { @@ -602,7 +487,6 @@ impl TryFrom for BatchPut { } Ok(BatchPut { - cluster_id: header.map_or(0, |h| h.cluster_id), kvs, options: Some(options), }) @@ -610,7 +494,6 @@ impl TryFrom for BatchPut { } struct BatchDelete { - cluster_id: u64, keys: Vec>, options: Option, } @@ -619,11 +502,7 @@ impl TryFrom for BatchDelete { type Error = error::Error; fn try_from(req: BatchDeleteRequest) -> Result { - let BatchDeleteRequest { - header, - keys, - prev_kv, - } = req; + let BatchDeleteRequest { keys, prev_kv } = req; let mut options = DeleteOptions::default(); if prev_kv { @@ -631,7 +510,6 @@ impl TryFrom for BatchDelete { } Ok(BatchDelete { - cluster_id: header.map_or(0, |h| h.cluster_id), keys, options: Some(options), }) @@ -639,7 +517,6 @@ impl TryFrom for BatchDelete { } struct CompareAndPut { - cluster_id: u64, key: Vec, expect: Vec, value: Vec, @@ -650,15 +527,9 @@ impl TryFrom for CompareAndPut { type Error = error::Error; fn try_from(req: CompareAndPutRequest) -> Result { - let CompareAndPutRequest { - header, - key, - expect, - value, - } = req; + let CompareAndPutRequest { key, expect, value } = req; Ok(CompareAndPut { - cluster_id: header.map_or(0, |h| h.cluster_id), key, expect, value, @@ -668,7 +539,6 @@ impl TryFrom for CompareAndPut { } struct Delete { - cluster_id: u64, key: Vec, options: Option, } @@ -678,7 +548,6 @@ impl TryFrom for Delete { fn try_from(req: DeleteRangeRequest) -> Result { let DeleteRangeRequest { - header, key, range_end, prev_kv, @@ -695,7 +564,6 @@ impl TryFrom for Delete { } Ok(Delete { - cluster_id: header.map_or(0, |h| h.cluster_id), key, options: Some(options), }) @@ -703,7 +571,6 @@ impl TryFrom for Delete { } struct MoveValue { - cluster_id: u64, from_key: Vec, to_key: Vec, delete_options: Option, @@ -713,14 +580,9 @@ impl TryFrom for MoveValue { type Error = error::Error; fn try_from(req: MoveValueRequest) -> Result { - let MoveValueRequest { - header, - from_key, - to_key, - } = req; + let MoveValueRequest { from_key, to_key } = req; Ok(MoveValue { - cluster_id: header.map_or(0, |h| h.cluster_id), from_key, to_key, delete_options: Some(DeleteOptions::default().with_prev_key()), @@ -739,7 +601,6 @@ mod tests { range_end: b"test_range_end".to_vec(), limit: 64, keys_only: true, - ..Default::default() }; let get: Get = req.try_into().unwrap(); @@ -754,7 +615,6 @@ mod tests { key: b"test_key".to_vec(), value: b"test_value".to_vec(), prev_kv: true, - ..Default::default() }; let put: Put = req.try_into().unwrap(); @@ -768,7 +628,6 @@ mod tests { fn test_parse_batch_get() { let req = BatchGetRequest { keys: vec![b"k1".to_vec(), b"k2".to_vec(), b"k3".to_vec()], - ..Default::default() }; let batch_get: BatchGet = req.try_into().unwrap(); @@ -787,13 +646,13 @@ mod tests { value: b"test_value".to_vec(), }], prev_kv: true, - ..Default::default() }; let batch_put: BatchPut = req.try_into().unwrap(); - assert_eq!(b"test_key".to_vec(), batch_put.kvs.get(0).unwrap().key); - assert_eq!(b"test_value".to_vec(), batch_put.kvs.get(0).unwrap().value); + let kv = batch_put.kvs.get(0).unwrap(); + assert_eq!(b"test_key", kv.key()); + assert_eq!(b"test_value", kv.value()); let _ = batch_put.options.unwrap(); } @@ -802,7 +661,6 @@ mod tests { let req = BatchDeleteRequest { keys: vec![b"k1".to_vec(), b"k2".to_vec(), b"k3".to_vec()], prev_kv: true, - ..Default::default() }; let batch_delete: BatchDelete = req.try_into().unwrap(); @@ -820,7 +678,6 @@ mod tests { key: b"test_key".to_vec(), expect: b"test_expect".to_vec(), value: b"test_value".to_vec(), - ..Default::default() }; let compare_and_put: CompareAndPut = req.try_into().unwrap(); @@ -837,7 +694,6 @@ mod tests { key: b"test_key".to_vec(), range_end: b"test_range_end".to_vec(), prev_kv: true, - ..Default::default() }; let delete: Delete = req.try_into().unwrap(); @@ -851,7 +707,6 @@ mod tests { let req = MoveValueRequest { from_key: b"test_from_key".to_vec(), to_key: b"test_to_key".to_vec(), - ..Default::default() }; let move_value: MoveValue = req.try_into().unwrap(); diff --git a/src/meta-srv/src/service/store/etcd_util.rs b/src/meta-srv/src/service/store/etcd_util.rs index e97e8edba3..5baaa8ae67 100644 --- a/src/meta-srv/src/service/store/etcd_util.rs +++ b/src/meta-srv/src/service/store/etcd_util.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::KeyValue; +use common_meta::rpc::KeyValue; pub struct KvPair<'a>(&'a etcd_client::KeyValue); diff --git a/src/meta-srv/src/service/store/ext.rs b/src/meta-srv/src/service/store/ext.rs deleted file mode 100644 index 85704780da..0000000000 --- a/src/meta-srv/src/service/store/ext.rs +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use api::v1::meta::{DeleteRangeRequest, KeyValue, RangeRequest}; - -use crate::error::Result; -use crate::service::store::kv::KvStore; - -#[async_trait::async_trait] -pub trait KvStoreExt { - /// Get the value by the given key. - async fn get(&self, key: Vec) -> Result>; - - /// Check if a key exists, it does not return the value. - async fn exists(&self, key: Vec) -> Result; - - /// Delete the value by the given key. If prev_kv is true, - /// the previous key-value pairs will be returned. - async fn delete(&self, key: Vec, prev_kv: bool) -> Result>; -} - -#[async_trait::async_trait] -impl KvStoreExt for T -where - T: KvStore + ?Sized, -{ - async fn get(&self, key: Vec) -> Result> { - let req = RangeRequest { - key, - ..Default::default() - }; - - let mut kvs = self.range(req).await?.kvs; - - Ok(kvs.pop()) - } - - async fn exists(&self, key: Vec) -> Result { - let req = RangeRequest { - key, - keys_only: true, - ..Default::default() - }; - - let kvs = self.range(req).await?.kvs; - - Ok(!kvs.is_empty()) - } - - async fn delete(&self, key: Vec, prev_kv: bool) -> Result> { - let req = DeleteRangeRequest { - key, - prev_kv, - ..Default::default() - }; - - let mut prev_kvs = self.delete_range(req).await?.prev_kvs; - - Ok(prev_kvs.pop()) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use api::v1::meta::PutRequest; - - use crate::service::store::ext::KvStoreExt; - use crate::service::store::kv::KvStoreRef; - use crate::service::store::memory::MemStore; - - #[tokio::test] - async fn test_get() { - let mut in_mem = Arc::new(MemStore::new()) as KvStoreRef; - - put_stats_to_store(&mut in_mem).await; - - let kv = in_mem - .get("test_key1".as_bytes().to_vec()) - .await - .unwrap() - .unwrap(); - - assert_eq!("test_key1".as_bytes(), kv.key); - assert_eq!("test_val1".as_bytes(), kv.value); - - let kv = in_mem - .get("test_key2".as_bytes().to_vec()) - .await - .unwrap() - .unwrap(); - - assert_eq!("test_key2".as_bytes(), kv.key); - assert_eq!("test_val2".as_bytes(), kv.value); - - let may_kv = in_mem.get("test_key3".as_bytes().to_vec()).await.unwrap(); - - assert!(may_kv.is_none()); - } - - #[tokio::test] - async fn test_exists() { - let mut in_mem = Arc::new(MemStore::new()) as KvStoreRef; - - put_stats_to_store(&mut in_mem).await; - - assert!(in_mem - .exists("test_key1".as_bytes().to_vec()) - .await - .unwrap()); - assert!(in_mem - .exists("test_key2".as_bytes().to_vec()) - .await - .unwrap()); - assert!(!in_mem - .exists("test_key3".as_bytes().to_vec()) - .await - .unwrap()); - assert!(!in_mem.exists("test_key".as_bytes().to_vec()).await.unwrap()); - } - - #[tokio::test] - async fn test_delete() { - let mut in_mem = Arc::new(MemStore::new()) as KvStoreRef; - - let mut prev_kv = in_mem - .delete("test_key1".as_bytes().to_vec(), true) - .await - .unwrap(); - assert!(prev_kv.is_none()); - - put_stats_to_store(&mut in_mem).await; - - assert!(in_mem - .exists("test_key1".as_bytes().to_vec()) - .await - .unwrap()); - - prev_kv = in_mem - .delete("test_key1".as_bytes().to_vec(), true) - .await - .unwrap(); - assert_eq!("test_key1".as_bytes(), prev_kv.unwrap().key); - } - - async fn put_stats_to_store(store: &mut KvStoreRef) { - assert!(store - .put(PutRequest { - key: "test_key1".as_bytes().to_vec(), - value: "test_val1".as_bytes().to_vec(), - ..Default::default() - }) - .await - .is_ok()); - - assert!(store - .put(PutRequest { - key: "test_key2".as_bytes().to_vec(), - value: "test_val2".as_bytes().to_vec(), - ..Default::default() - }) - .await - .is_ok()); - } -} diff --git a/src/meta-srv/src/service/store/kv.rs b/src/meta-srv/src/service/store/kv.rs index 5c03e9362e..febbf3a7fa 100644 --- a/src/meta-srv/src/service/store/kv.rs +++ b/src/meta-srv/src/service/store/kv.rs @@ -14,38 +14,13 @@ use std::sync::Arc; -use api::v1::meta::{ - BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, - BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, - DeleteRangeResponse, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, - RangeRequest, RangeResponse, -}; +use common_meta::kv_backend::KvBackend; -use crate::error::Result; -use crate::service::store::txn::TxnService; +use crate::error::Error; -pub type KvStoreRef = Arc; +pub type KvStoreRef = Arc>; pub type ResettableKvStoreRef = Arc; -#[async_trait::async_trait] -pub trait KvStore: TxnService { - async fn range(&self, req: RangeRequest) -> Result; - - async fn put(&self, req: PutRequest) -> Result; - - async fn batch_get(&self, req: BatchGetRequest) -> Result; - - async fn batch_put(&self, req: BatchPutRequest) -> Result; - - async fn batch_delete(&self, req: BatchDeleteRequest) -> Result; - - async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result; - - async fn delete_range(&self, req: DeleteRangeRequest) -> Result; - - async fn move_value(&self, req: MoveValueRequest) -> Result; -} - -pub trait ResettableKvStore: KvStore { +pub trait ResettableKvStore: KvBackend { fn reset(&self); } diff --git a/src/meta-srv/src/service/store/memory.rs b/src/meta-srv/src/service/store/memory.rs index 0d9e58f776..51ded2ab3b 100644 --- a/src/meta-srv/src/service/store/memory.rs +++ b/src/meta-srv/src/service/store/memory.rs @@ -12,697 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::btree_map::Entry; -use std::collections::BTreeMap; -use std::ops::Range; +use common_meta::kv_backend::memory::MemoryKvBackend; -use api::v1::meta::{ - BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest, - BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest, - DeleteRangeResponse, KeyValue, MoveValueRequest, MoveValueResponse, PutRequest, PutResponse, - RangeRequest, RangeResponse, ResponseHeader, -}; -use common_telemetry::timer; -use parking_lot::RwLock; +use crate::error::Error; +use crate::service::store::kv::ResettableKvStore; -use super::ext::KvStoreExt; -use crate::error::Result; -use crate::metrics::{METRIC_META_KV_REQUEST, METRIC_META_TXN_REQUEST}; -use crate::service::store::kv::{KvStore, ResettableKvStore}; -use crate::service::store::txn::{Txn, TxnOp, TxnOpResponse, TxnRequest, TxnResponse, TxnService}; - -pub struct MemStore { - inner: RwLock, Vec>>, -} - -impl Default for MemStore { - fn default() -> Self { - Self::new() - } -} - -impl MemStore { - pub fn new() -> Self { - Self { - inner: RwLock::new(Default::default()), - } - } -} +pub type MemStore = MemoryKvBackend; impl ResettableKvStore for MemStore { fn reset(&self) { - self.inner.write().clear(); - } -} - -#[async_trait::async_trait] -impl KvStore for MemStore { - async fn range(&self, req: RangeRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "range"),] - ); - - let RangeRequest { - header, - key, - range_end, - limit, - keys_only, - } = req; - - let memory = self.inner.read(); - - let mut kvs = if range_end.is_empty() { - memory.get_key_value(&key).map_or(vec![], |(k, v)| { - vec![KeyValue { - key: k.clone(), - value: if keys_only { vec![] } else { v.clone() }, - }] - }) - } else { - let range = Range { - start: key, - end: range_end, - }; - memory - .range(range) - .map(|kv| KeyValue { - key: kv.0.clone(), - value: if keys_only { vec![] } else { kv.1.clone() }, - }) - .collect::>() - }; - - let more = if limit > 0 { - kvs.truncate(limit as usize); - true - } else { - false - }; - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(RangeResponse { header, kvs, more }) - } - - async fn put(&self, req: PutRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "put"),] - ); - - let PutRequest { - header, - key, - value, - prev_kv, - } = req; - - let mut memory = self.inner.write(); - - let prev_value = memory.insert(key.clone(), value); - let prev_kv = if prev_kv { - prev_value.map(|value| KeyValue { key, value }) - } else { - None - }; - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(PutResponse { header, prev_kv }) - } - - async fn batch_get(&self, req: BatchGetRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "batch_get"),] - ); - - let BatchGetRequest { header, keys } = req; - - let mut kvs = Vec::with_capacity(keys.len()); - for key in keys { - if let Some(kv) = self.get(key).await? { - kvs.push(kv); - } - } - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(BatchGetResponse { header, kvs }) - } - - async fn batch_put(&self, req: BatchPutRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "batch_put"),] - ); - - let BatchPutRequest { - header, - kvs, - prev_kv, - } = req; - - let mut memory = self.inner.write(); - - let prev_kvs = if prev_kv { - kvs.into_iter() - .map(|kv| (kv.key.clone(), memory.insert(kv.key, kv.value))) - .filter(|(_, v)| v.is_some()) - .map(|(key, value)| KeyValue { - key, - value: value.unwrap(), - }) - .collect() - } else { - for kv in kvs.into_iter() { - let _ = memory.insert(kv.key, kv.value); - } - vec![] - }; - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(BatchPutResponse { header, prev_kvs }) - } - - async fn batch_delete(&self, req: BatchDeleteRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "batch_delete"),] - ); - - let BatchDeleteRequest { - header, - keys, - prev_kv, - } = req; - - let mut memory = self.inner.write(); - - let prev_kvs = if prev_kv { - keys.into_iter() - .filter_map(|key| memory.remove(&key).map(|value| KeyValue { key, value })) - .collect() - } else { - for key in keys.into_iter() { - let _ = memory.remove(&key); - } - vec![] - }; - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(BatchDeleteResponse { header, prev_kvs }) - } - - async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "compare_and_put"),] - ); - - let CompareAndPutRequest { - header, - key, - expect, - value, - } = req; - - let mut memory = self.inner.write(); - - let (success, prev_kv) = match memory.entry(key) { - Entry::Vacant(e) => { - let success = expect.is_empty(); - if success { - let _ = e.insert(value); - } - (success, None) - } - Entry::Occupied(mut e) => { - let key = e.key().clone(); - let prev_val = e.get().clone(); - let success = prev_val == expect; - if success { - let _ = e.insert(value); - } - (success, Some((key, prev_val))) - } - }; - - let prev_kv = prev_kv.map(|(key, value)| KeyValue { key, value }); - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(CompareAndPutResponse { - header, - success, - prev_kv, - }) - } - - async fn delete_range(&self, req: DeleteRangeRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "deleteL_range"),] - ); - - let DeleteRangeRequest { - header, - key, - range_end, - prev_kv, - } = req; - - let mut memory = self.inner.write(); - - let prev_kvs = if range_end.is_empty() { - let prev_val = memory.remove(&key); - prev_val.map_or(vec![], |value| vec![KeyValue { key, value }]) - } else { - let range = Range { - start: key, - end: range_end, - }; - memory - .drain_filter(|key, _| range.contains(key)) - .map(|(key, value)| KeyValue { key, value }) - .collect::>() - }; - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(DeleteRangeResponse { - header, - deleted: prev_kvs.len() as i64, - prev_kvs: if prev_kv { - prev_kvs - } else { - Default::default() - }, - }) - } - - async fn move_value(&self, req: MoveValueRequest) -> Result { - let _timer = timer!( - METRIC_META_KV_REQUEST, - &[("target", "memory"), ("op", "move_value"),] - ); - - let MoveValueRequest { - header, - from_key, - to_key, - } = req; - - let mut memory = self.inner.write(); - - let kv = match memory.remove(&from_key) { - Some(v) => { - let _ = memory.insert(to_key, v.clone()); - Some((from_key, v)) - } - None => memory.get(&to_key).map(|v| (to_key, v.clone())), - }; - - let kv = kv.map(|(key, value)| KeyValue { key, value }); - - let cluster_id = header.map_or(0, |h| h.cluster_id); - let header = Some(ResponseHeader::success(cluster_id)); - Ok(MoveValueResponse { header, kv }) - } -} - -#[async_trait::async_trait] -impl TxnService for MemStore { - async fn txn(&self, txn: Txn) -> Result { - let _timer = timer!( - METRIC_META_TXN_REQUEST, - &[("target", "memory".to_string()), ("op", "txn".to_string()),] - ); - - let TxnRequest { - compare, - success, - failure, - } = txn.into(); - - let mut memory = self.inner.write(); - - let succeeded = compare - .iter() - .all(|x| x.compare_with_value(memory.get(&x.key))); - - let do_txn = |txn_op| match txn_op { - TxnOp::Put(key, value) => { - let prev_value = memory.insert(key.clone(), value); - let prev_kv = prev_value.map(|value| KeyValue { key, value }); - let put_res = PutResponse { - prev_kv, - ..Default::default() - }; - TxnOpResponse::ResponsePut(put_res) - } - TxnOp::Get(key) => { - let value = memory.get(&key); - let kv = value.map(|value| KeyValue { - key, - value: value.clone(), - }); - let get_res = RangeResponse { - kvs: kv.map(|kv| vec![kv]).unwrap_or(vec![]), - ..Default::default() - }; - TxnOpResponse::ResponseGet(get_res) - } - TxnOp::Delete(key) => { - let prev_value = memory.remove(&key); - let prev_kv = prev_value.map(|value| KeyValue { key, value }); - let delete_res = DeleteRangeResponse { - prev_kvs: prev_kv.map(|kv| vec![kv]).unwrap_or(vec![]), - ..Default::default() - }; - TxnOpResponse::ResponseDelete(delete_res) - } - }; - - let responses: Vec<_> = if succeeded { - success.into_iter().map(do_txn).collect() - } else { - failure.into_iter().map(do_txn).collect() - }; - - Ok(TxnResponse { - succeeded, - responses, - }) - } -} - -#[cfg(test)] -mod tests { - use std::sync::atomic::{AtomicU8, Ordering}; - use std::sync::Arc; - - use api::v1::meta::{ - BatchGetRequest, BatchPutRequest, CompareAndPutRequest, DeleteRangeRequest, KeyValue, - MoveValueRequest, PutRequest, RangeRequest, - }; - - use super::MemStore; - use crate::service::store::ext::KvStoreExt; - use crate::service::store::kv::KvStore; - use crate::util; - - async fn mock_mem_store_with_data() -> MemStore { - let kv_store = MemStore::new(); - let kvs = mock_kvs(); - - assert!(kv_store - .batch_put(BatchPutRequest { - kvs, - ..Default::default() - }) - .await - .is_ok()); - - assert!(kv_store - .put(PutRequest { - key: b"key11".to_vec(), - value: b"val11".to_vec(), - ..Default::default() - }) - .await - .is_ok()); - - kv_store - } - - fn mock_kvs() -> Vec { - vec![ - KeyValue { - key: b"key1".to_vec(), - value: b"val1".to_vec(), - }, - KeyValue { - key: b"key2".to_vec(), - value: b"val2".to_vec(), - }, - KeyValue { - key: b"key3".to_vec(), - value: b"val3".to_vec(), - }, - ] - } - - #[tokio::test] - async fn test_put() { - let kv_store = mock_mem_store_with_data().await; - - let resp = kv_store - .put(PutRequest { - key: b"key11".to_vec(), - value: b"val12".to_vec(), - ..Default::default() - }) - .await - .unwrap(); - assert!(resp.prev_kv.is_none()); - - let resp = kv_store - .put(PutRequest { - key: b"key11".to_vec(), - value: b"val13".to_vec(), - prev_kv: true, - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(b"key11".as_slice(), resp.prev_kv.as_ref().unwrap().key); - assert_eq!(b"val12".as_slice(), resp.prev_kv.as_ref().unwrap().value); - } - - #[tokio::test] - async fn test_range() { - let kv_store = mock_mem_store_with_data().await; - - let key = b"key1".to_vec(); - let range_end = util::get_prefix_end_key(b"key1"); - - let resp = kv_store - .range(RangeRequest { - key: key.clone(), - range_end: range_end.clone(), - limit: 0, - keys_only: false, - ..Default::default() - }) - .await - .unwrap(); - - assert_eq!(2, resp.kvs.len()); - assert_eq!(b"key1".as_slice(), resp.kvs[0].key); - assert_eq!(b"val1".as_slice(), resp.kvs[0].value); - assert_eq!(b"key11".as_slice(), resp.kvs[1].key); - assert_eq!(b"val11".as_slice(), resp.kvs[1].value); - - let resp = kv_store - .range(RangeRequest { - key: key.clone(), - range_end: range_end.clone(), - limit: 0, - keys_only: true, - ..Default::default() - }) - .await - .unwrap(); - - assert_eq!(2, resp.kvs.len()); - assert_eq!(b"key1".as_slice(), resp.kvs[0].key); - assert_eq!(b"".as_slice(), resp.kvs[0].value); - assert_eq!(b"key11".as_slice(), resp.kvs[1].key); - assert_eq!(b"".as_slice(), resp.kvs[1].value); - - let resp = kv_store - .range(RangeRequest { - key: key.clone(), - limit: 0, - keys_only: false, - ..Default::default() - }) - .await - .unwrap(); - - assert_eq!(1, resp.kvs.len()); - assert_eq!(b"key1".as_slice(), resp.kvs[0].key); - assert_eq!(b"val1".as_slice(), resp.kvs[0].value); - - let resp = kv_store - .range(RangeRequest { - key, - range_end, - limit: 1, - keys_only: false, - ..Default::default() - }) - .await - .unwrap(); - - assert_eq!(1, resp.kvs.len()); - assert_eq!(b"key1".as_slice(), resp.kvs[0].key); - assert_eq!(b"val1".as_slice(), resp.kvs[0].value); - } - - #[tokio::test] - async fn test_batch_get() { - let kv_store = mock_mem_store_with_data().await; - - let keys = vec![]; - let batch_resp = kv_store - .batch_get(BatchGetRequest { - keys, - ..Default::default() - }) - .await - .unwrap(); - - assert!(batch_resp.kvs.is_empty()); - - let keys = vec![b"key10".to_vec()]; - let batch_resp = kv_store - .batch_get(BatchGetRequest { - keys, - ..Default::default() - }) - .await - .unwrap(); - - assert!(batch_resp.kvs.is_empty()); - - let keys = vec![b"key1".to_vec(), b"key3".to_vec(), b"key4".to_vec()]; - let batch_resp = kv_store - .batch_get(BatchGetRequest { - keys, - ..Default::default() - }) - .await - .unwrap(); - - assert_eq!(2, batch_resp.kvs.len()); - assert_eq!(b"key1".as_slice(), batch_resp.kvs[0].key); - assert_eq!(b"val1".as_slice(), batch_resp.kvs[0].value); - assert_eq!(b"key3".as_slice(), batch_resp.kvs[1].key); - assert_eq!(b"val3".as_slice(), batch_resp.kvs[1].value); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_compare_and_put() { - let kv_store = Arc::new(MemStore::new()); - let success = Arc::new(AtomicU8::new(0)); - - let mut joins = vec![]; - for _ in 0..20 { - let kv_store_clone = kv_store.clone(); - let success_clone = success.clone(); - let join = tokio::spawn(async move { - let req = CompareAndPutRequest { - key: b"key".to_vec(), - expect: vec![], - value: b"val_new".to_vec(), - ..Default::default() - }; - let resp = kv_store_clone.compare_and_put(req).await.unwrap(); - if resp.success { - let _ = success_clone.fetch_add(1, Ordering::SeqCst); - } - }); - joins.push(join); - } - - for join in joins { - join.await.unwrap(); - } - - assert_eq!(1, success.load(Ordering::SeqCst)); - } - - #[tokio::test] - async fn test_delete_range() { - let kv_store = mock_mem_store_with_data().await; - - let req = DeleteRangeRequest { - key: b"key3".to_vec(), - range_end: vec![], - prev_kv: true, - ..Default::default() - }; - - let resp = kv_store.delete_range(req).await.unwrap(); - assert_eq!(1, resp.prev_kvs.len()); - assert_eq!(b"key3".as_slice(), resp.prev_kvs[0].key); - assert_eq!(b"val3".as_slice(), resp.prev_kvs[0].value); - - let get_resp = kv_store.get(b"key3".to_vec()).await.unwrap(); - assert!(get_resp.is_none()); - - let req = DeleteRangeRequest { - key: b"key2".to_vec(), - range_end: vec![], - prev_kv: false, - ..Default::default() - }; - - let resp = kv_store.delete_range(req).await.unwrap(); - assert!(resp.prev_kvs.is_empty()); - - let get_resp = kv_store.get(b"key2".to_vec()).await.unwrap(); - assert!(get_resp.is_none()); - - let key = b"key1".to_vec(); - let range_end = util::get_prefix_end_key(b"key1"); - - let req = DeleteRangeRequest { - key: key.clone(), - range_end: range_end.clone(), - prev_kv: true, - ..Default::default() - }; - let resp = kv_store.delete_range(req).await.unwrap(); - assert_eq!(2, resp.prev_kvs.len()); - - let req = RangeRequest { - key, - range_end, - ..Default::default() - }; - let resp = kv_store.range(req).await.unwrap(); - assert!(resp.kvs.is_empty()); - } - - #[tokio::test] - async fn test_move_value() { - let kv_store = mock_mem_store_with_data().await; - - let req = MoveValueRequest { - from_key: b"key1".to_vec(), - to_key: b"key111".to_vec(), - ..Default::default() - }; - - let resp = kv_store.move_value(req).await.unwrap(); - assert_eq!(b"key1".as_slice(), resp.kv.as_ref().unwrap().key); - assert_eq!(b"val1".as_slice(), resp.kv.as_ref().unwrap().value); - - let kv_store = mock_mem_store_with_data().await; - - let req = MoveValueRequest { - from_key: b"notexistkey".to_vec(), - to_key: b"key222".to_vec(), - ..Default::default() - }; - - let resp = kv_store.move_value(req).await.unwrap(); - assert!(resp.kv.is_none()); + self.clear(); } } diff --git a/src/meta-srv/src/table_routes.rs b/src/meta-srv/src/table_routes.rs index 111ad30014..8287bdf680 100644 --- a/src/meta-srv/src/table_routes.rs +++ b/src/meta-srv/src/table_routes.rs @@ -14,26 +14,25 @@ use std::collections::HashMap; -use api::v1::meta::{MoveValueRequest, PutRequest, TableRouteValue}; +use api::v1::meta::TableRouteValue; use catalog::helper::{TableGlobalKey, TableGlobalValue}; use common_meta::key::TableRouteKey; -use common_meta::rpc::store::{BatchGetRequest, BatchGetResponse}; +use common_meta::rpc::store::{BatchGetRequest, MoveValueRequest, PutRequest}; use common_telemetry::warn; use snafu::{OptionExt, ResultExt}; use table::engine::TableReference; use crate::error::{ - ConvertProtoDataSnafu, DecodeTableRouteSnafu, InvalidCatalogValueSnafu, Result, - TableNotFoundSnafu, TableRouteNotFoundSnafu, + DecodeTableRouteSnafu, InvalidCatalogValueSnafu, Result, TableNotFoundSnafu, + TableRouteNotFoundSnafu, }; -use crate::service::store::ext::KvStoreExt; use crate::service::store::kv::KvStoreRef; pub async fn get_table_global_value( kv_store: &KvStoreRef, key: &TableGlobalKey, ) -> Result> { - let kv = kv_store.get(key.to_raw_key()).await?; + let kv = kv_store.get(&key.to_raw_key()).await?; kv.map(|kv| TableGlobalValue::from_bytes(kv.value).context(InvalidCatalogValueSnafu)) .transpose() } @@ -45,13 +44,8 @@ pub(crate) async fn batch_get_table_global_value( let req = BatchGetRequest { keys: keys.iter().map(|x| x.to_raw_key()).collect::>(), }; - let mut resp: BatchGetResponse = kv_store - .batch_get(req.into()) - .await? - .try_into() - .context(ConvertProtoDataSnafu)?; + let kvs = kv_store.batch_get(req).await?.kvs; - let kvs = resp.take_kvs(); let mut result = HashMap::with_capacity(kvs.len()); for kv in kvs { let key = TableGlobalKey::try_from_raw_key(kv.key()).context(InvalidCatalogValueSnafu)?; @@ -73,7 +67,6 @@ pub(crate) async fn put_table_global_value( value: &TableGlobalValue, ) -> Result<()> { let req = PutRequest { - header: None, key: key.to_raw_key(), value: value.as_bytes().context(InvalidCatalogValueSnafu)?, prev_kv: false, @@ -101,15 +94,12 @@ pub(crate) async fn get_table_route_value( key: &TableRouteKey<'_>, ) -> Result { let kv = kv_store - .get(key.to_string().into_bytes()) + .get(key.to_string().as_bytes()) .await? .with_context(|| TableRouteNotFoundSnafu { key: key.to_string(), })?; - kv.value - .as_slice() - .try_into() - .context(DecodeTableRouteSnafu) + kv.value().try_into().context(DecodeTableRouteSnafu) } pub(crate) async fn put_table_route_value( @@ -118,7 +108,6 @@ pub(crate) async fn put_table_route_value( value: TableRouteValue, ) -> Result<()> { let req = PutRequest { - header: None, key: key.to_string().into_bytes(), value: value.into(), prev_kv: false, @@ -150,14 +139,10 @@ async fn move_value( ) -> Result, Vec)>> { let from_key = from_key.into(); let to_key = to_key.into(); - let move_req = MoveValueRequest { - from_key, - to_key, - ..Default::default() - }; + let move_req = MoveValueRequest { from_key, to_key }; let res = kv_store.move_value(move_req).await?; - Ok(res.kv.map(|kv| (kv.key, kv.value))) + Ok(res.0.map(Into::into)) } pub(crate) fn table_route_key(table_id: u32, t: &TableGlobalKey) -> TableRouteKey<'_> { diff --git a/tests-integration/tests/region_failover.rs b/tests-integration/tests/region_failover.rs index e4a1eb6463..02f57e6a56 100644 --- a/tests-integration/tests/region_failover.rs +++ b/tests-integration/tests/region_failover.rs @@ -21,8 +21,8 @@ use catalog::helper::TableGlobalKey; use catalog::remote::CachedMetaKvBackend; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, MITO_ENGINE}; use common_meta::ident::TableIdent; -use common_meta::kv_backend::Kv; use common_meta::rpc::router::TableRoute; +use common_meta::rpc::KeyValue; use common_meta::table_name::TableName; use common_meta::RegionIdent; use common_procedure::{watcher, ProcedureWithId}; @@ -175,7 +175,7 @@ pub async fn test_region_failover(store_type: StorageType) { assert!(success) } -fn get_table_cache(instance: &Arc, key: &str) -> Option> { +fn get_table_cache(instance: &Arc, key: &str) -> Option> { let catalog_manager = instance .catalog_manager() .as_any()