refactor: unify KvBackend and KvStore (#1890)

* refactor: unify KvBackend and KvStore
This commit is contained in:
LFC
2023-07-07 19:53:49 +08:00
committed by GitHub
parent 4cc42e2ba6
commit f19498f73e
54 changed files with 2116 additions and 2156 deletions

1
Cargo.lock generated
View File

@@ -1812,6 +1812,7 @@ dependencies = [
"common-telemetry",
"common-time",
"datatypes",
"etcd-client",
"futures",
"lazy_static",
"prost",

View File

@@ -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" }

View File

@@ -30,76 +30,3 @@ pub trait KvCacheInvalidator: Send + Sync {
}
pub type KvCacheInvalidatorRef = Arc<dyn KvCacheInvalidator>;
#[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<Result<(), Option<Vec<u8>>>, 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());
}
}

View File

@@ -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<Cache<Vec<u8>, Kv>>;
pub type CacheBackendRef = Arc<Cache<Vec<u8>, 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<Option<Kv>> {
async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
self.kv_backend.range(req).await
}
async fn get(&self, key: &[u8]) -> Result<Option<KeyValue>> {
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<PutResponse> {
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<BatchPutResponse> {
let keys = req
.kvs
.iter()
.map(|kv| kv.key().to_vec())
.collect::<Vec<_>>();
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<DeleteRangeResponse> {
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<BatchDeleteResponse> {
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<BatchGetResponse> {
self.kv_backend.batch_get(req).await
}
async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result<CompareAndPutResponse> {
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<MoveValueResponse> {
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<std::result::Result<(), Option<Vec<u8>>>> {
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<MetaClient>) -> 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<MetaClient>,
}
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<Option<Kv>> {
async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
self.client
.range(req)
.await
.map_err(BoxedError::new)
.context(MetaSrvSnafu)
}
async fn get(&self, key: &[u8]) -> Result<Option<KeyValue>> {
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<BatchPutResponse> {
self.client
.batch_put(req)
.await
.map_err(BoxedError::new)
.context(MetaSrvSnafu)
}
async fn put(&self, req: PutRequest) -> Result<PutResponse> {
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<DeleteRangeResponse> {
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<BatchDeleteResponse> {
self.client
.batch_delete(req)
.await
.map_err(BoxedError::new)
.context(MetaSrvSnafu)
}
async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse> {
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<std::result::Result<(), Option<Vec<u8>>>> {
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<CompareAndPutResponse> {
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<MoveValueResponse> {
self.client
.move_value(req)
.await
.map_err(BoxedError::new)
.context(MetaSrvSnafu)?;
Ok(())
.context(MetaSrvSnafu)
}
fn as_any(&self) -> &dyn Any {

View File

@@ -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<Box<dyn Stream<Item = Result<CatalogKey>> + Send + '_>> {
async fn iter_remote_catalogs(&self) -> Result<Vec<CatalogKey>> {
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<u32> {
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::<Vec<_>>().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<Box<dyn Stream<Item = Result<SchemaKey>> + Send + 'a>> {
) -> Result<Vec<SchemaKey>> {
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<Box<dyn Stream<Item = Result<(TableGlobalKey, TableGlobalValue)>> + Send + 'a>> {
) -> Result<Vec<(TableGlobalKey, TableGlobalValue)>> {
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<Vec<String>> {
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<Vec<String>> {
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<bool> {
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);

View File

@@ -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::<Error>::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::<Vec<_>>()
@@ -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));

View File

@@ -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

View File

@@ -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<T> = std::result::Result<T, Error>;
@@ -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 { .. }

View File

@@ -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<bool> {
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(&region),
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(&region),
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<Vec<DatanodeTableValue>> {
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::<Vec<_>>()
.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::<Result<Vec<_>>>()?;
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();

View File

@@ -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<u8>))`
/// will be returned, the `Err(Vec<u8>)` 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<TableInfoValue>,
@@ -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);
}

View File

@@ -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<Option<TableNameValue>> {
@@ -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<Vec<String>> {
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::<Vec<_>>()
.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::<Result<Vec<_>>>()?;
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();

View File

@@ -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<DatanodeId, Vec<RegionNumber>>;
@@ -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<u8>))`
/// will be returned, the `Err(Vec<u8>)` 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<TableRegionValue>,
@@ -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);
}

View File

@@ -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<u8>, pub Vec<u8>);
pub type ValueIter<'a, E> = Pin<Box<dyn Stream<Item = Result<Kv, E>> + Send + 'a>>;
pub type KvBackendRef = Arc<dyn KvBackend<Error = Error>>;
pub type KvBackendRef = Arc<dyn KvBackend<Error = Error> + 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<RangeResponse, Self::Error>;
async fn set(&self, key: &[u8], val: &[u8]) -> Result<(), Self::Error>;
async fn put(&self, req: PutRequest) -> Result<PutResponse, Self::Error>;
/// 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<u8>))`
/// will be returned, the `Err(Vec<u8>)` 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<BatchPutResponse, Self::Error>;
async fn compare_and_put(
&self,
key: &[u8],
expect: &[u8],
val: &[u8],
) -> Result<Result<(), Option<Vec<u8>>>, Self::Error>;
req: CompareAndPutRequest,
) -> Result<CompareAndPutResponse, Self::Error>;
async fn delete_range(&self, key: &[u8], end: &[u8]) -> Result<(), Self::Error>;
async fn delete_range(
&self,
req: DeleteRangeRequest,
) -> Result<DeleteRangeResponse, Self::Error>;
async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
self.delete_range(key, &[]).await
async fn delete(&self, key: &[u8], prev_kv: bool) -> Result<Option<KeyValue>, 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<BatchDeleteResponse, Self::Error>;
/// Default get is implemented based on `range` method.
async fn get(&self, key: &[u8]) -> Result<Option<Kv>, 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<Option<KeyValue>, 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<BatchGetResponse, Self::Error>;
/// 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<MoveValueResponse, Self::Error>;
fn as_any(&self) -> &dyn Any;
}

View File

@@ -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<T> {
kvs: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
_phantom: PhantomData<T>,
}
impl Display for MemoryKvBackend {
impl<T> Display for MemoryKvBackend<T> {
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<T> Default for MemoryKvBackend<T> {
fn default() -> Self {
Self {
kvs: RwLock::new(BTreeMap::new()),
_phantom: PhantomData,
}
}
}
impl<T> MemoryKvBackend<T> {
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<T: ErrorExt + Send + Sync + 'static> KvBackend for MemoryKvBackend<T> {
fn name(&self) -> &str {
"Memory"
}
async fn range(&self, req: RangeRequest) -> Result<RangeResponse, Self::Error> {
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<dyn Iterator<Item = (&Vec<u8>, &Vec<u8>)>> = 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::<Vec<_>>();
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<PutResponse, Self::Error> {
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<BatchPutResponse, Self::Error> {
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<Result<(), Option<Vec<u8>>>, Error> {
let key = key.to_vec();
let val = val.to_vec();
req: CompareAndPutRequest,
) -> Result<CompareAndPutResponse, Self::Error> {
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<DeleteRangeResponse, Self::Error> {
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::<Vec<_>>()
} else {
let range = Range {
start: key,
end: range_end,
};
kvs.drain_filter(|key, _| range.contains(key))
.map(Into::into)
.collect::<Vec<_>>()
};
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<BatchDeleteResponse, Self::Error> {
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<BatchGetResponse, Self::Error> {
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::<Vec<_>>();
Ok(BatchGetResponse { kvs })
}
async fn move_value(&self, req: MoveValueRequest) -> Result<MoveValueResponse, Self::Error> {
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<T: ErrorExt + Send + Sync> TxnService for MemoryKvBackend<T> {
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<TxnResponse, Self::Error> {
let _timer = timer!(
METRIC_META_TXN_REQUEST,
&[("target", "memory"), ("op", "txn")]
);
let range = backend.range(b"key").try_collect::<Vec<_>>().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<Error> {
let kv_store = MemoryKvBackend::<Error>::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<KeyValue> {
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::<Error>::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());
}
}

View File

@@ -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<TxnResponse> {
type Error: ErrorExt;
async fn txn(&self, _txn: Txn) -> Result<TxnResponse, Self::Error> {
unimplemented!("txn is not implemented")
}
}
@@ -169,11 +171,14 @@ impl From<Txn> 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::<Error>::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

View File

@@ -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<Txn> for EtcdTxn {
fn from(txn: Txn) -> Self {
@@ -88,31 +87,25 @@ impl TryFrom<EtcdTxnOpResponse> for TxnOpResponse {
fn try_from(op_resp: EtcdTxnOpResponse) -> Result<Self> {
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::<Vec<_>>();
let delete_res = DeleteRangeResponse {
prev_kvs,
deleted: res.deleted(),
..Default::default()
};
Ok(TxnOpResponse::ResponseDelete(delete_res))
}

View File

@@ -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;

View File

@@ -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";

View File

@@ -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<u8>,
pub value: Vec<u8>,
}
impl From<KeyValue> for PbKeyValue {
fn from(kv: KeyValue) -> Self {
Self {
key: kv.key,
value: kv.value,
}
}
}
impl From<etcd_client::KeyValue> for KeyValue {
fn from(kv: etcd_client::KeyValue) -> Self {
Self {
key: kv.key().to_vec(),
value: kv.value().to_vec(),
}
}
}
impl From<KeyValue> for (Vec<u8>, Vec<u8>) {
fn from(kv: KeyValue) -> Self {
(kv.key, kv.value)
}
}
impl From<(Vec<u8>, Vec<u8>)> for KeyValue {
fn from(kv: (Vec<u8>, Vec<u8>)) -> 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<u8> {
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<u8> {
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());

View File

@@ -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<RangeRequest> for PbRangeRequest {
}
}
impl From<PbRangeRequest> 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<KeyValue>,
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::<Vec<_>>()
.join(", "),
self.more
)
}
}
impl TryFrom<PbRangeResponse> for RangeResponse {
type Error = error::Error;
@@ -128,29 +172,25 @@ impl TryFrom<PbRangeResponse> for RangeResponse {
fn try_from(pb: PbRangeResponse) -> Result<Self> {
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<ResponseHeader> {
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<KeyValue> {
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<PutRequest> for PbPutRequest {
}
}
impl From<PbPutRequest> 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<KeyValue>,
}
impl TryFrom<PbPutResponse> for PutResponse {
type Error = error::Error;
@@ -220,27 +272,22 @@ impl TryFrom<PbPutResponse> for PutResponse {
fn try_from(pb: PbPutResponse) -> Result<Self> {
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<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn take_prev_kv(&mut self) -> Option<KeyValue> {
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<Vec<u8>>,
}
@@ -254,6 +301,12 @@ impl From<BatchGetRequest> for PbBatchGetRequest {
}
}
impl From<PbBatchGetRequest> 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<KeyValue>,
}
impl Display for BatchGetResponse {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"[{}]",
self.kvs
.iter()
.map(|kv| kv.to_string())
.collect::<Vec<_>>()
.join(", "),
)
}
}
impl TryFrom<PbBatchGetResponse> for BatchGetResponse {
type Error = error::Error;
@@ -282,30 +351,24 @@ impl TryFrom<PbBatchGetResponse> for BatchGetResponse {
fn try_from(pb: PbBatchGetResponse) -> Result<Self> {
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<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn take_kvs(&mut self) -> Vec<KeyValue> {
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<PbKeyValue>,
pub kvs: Vec<KeyValue>,
/// 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<BatchPutRequest> 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<PbBatchPutRequest> 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<Vec<u8>>, value: impl Into<Vec<u8>>) -> 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<KeyValue>,
}
impl TryFrom<PbBatchPutResponse> for BatchPutResponse {
type Error = error::Error;
@@ -357,24 +431,23 @@ impl TryFrom<PbBatchPutResponse> for BatchPutResponse {
fn try_from(pb: PbBatchPutResponse) -> Result<Self> {
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<ResponseHeader> {
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<KeyValue> {
self.0.prev_kvs.drain(..).map(KeyValue::new).collect()
self.prev_kvs.drain(..).collect()
}
}
@@ -396,6 +469,15 @@ impl From<BatchDeleteRequest> for PbBatchDeleteRequest {
}
}
impl From<PbBatchDeleteRequest> 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<KeyValue>,
}
impl TryFrom<PbBatchDeleteResponse> for BatchDeleteResponse {
type Error = error::Error;
@@ -429,24 +513,18 @@ impl TryFrom<PbBatchDeleteResponse> for BatchDeleteResponse {
fn try_from(pb: PbBatchDeleteResponse) -> Result<Self> {
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<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn take_prev_kvs(&mut self) -> Vec<KeyValue> {
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<CompareAndPutRequest> for PbCompareAndPutRequest {
}
}
impl From<PbCompareAndPutRequest> 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<KeyValue>,
}
impl TryFrom<PbCompareAndPutResponse> for CompareAndPutResponse {
type Error = error::Error;
@@ -513,29 +604,30 @@ impl TryFrom<PbCompareAndPutResponse> for CompareAndPutResponse {
fn try_from(pb: PbCompareAndPutResponse) -> Result<Self> {
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<ResponseHeader> {
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<KeyValue> {
self.0.prev_kv.take().map(KeyValue::new)
self.prev_kv.take()
}
}
@@ -571,6 +663,16 @@ impl From<DeleteRangeRequest> for PbDeleteRangeRequest {
}
}
impl From<PbDeleteRangeRequest> 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<KeyValue>,
}
impl TryFrom<PbDeleteRangeResponse> for DeleteRangeResponse {
type Error = error::Error;
@@ -633,29 +738,30 @@ impl TryFrom<PbDeleteRangeResponse> for DeleteRangeResponse {
fn try_from(pb: PbDeleteRangeResponse) -> Result<Self> {
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<ResponseHeader> {
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<KeyValue> {
self.0.prev_kvs.drain(..).map(KeyValue::new).collect()
self.prev_kvs.drain(..).collect()
}
}
@@ -678,6 +784,15 @@ impl From<MoveValueRequest> for PbMoveValueRequest {
}
}
impl From<PbMoveValueRequest> 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<Vec<u8>>, to_key: impl Into<Vec<u8>>) -> Self {
@@ -689,7 +804,7 @@ impl MoveValueRequest {
}
#[derive(Debug, Clone)]
pub struct MoveValueResponse(PbMoveValueResponse);
pub struct MoveValueResponse(pub Option<KeyValue>);
impl TryFrom<PbMoveValueResponse> for MoveValueResponse {
type Error = error::Error;
@@ -697,24 +812,21 @@ impl TryFrom<PbMoveValueResponse> for MoveValueResponse {
fn try_from(pb: PbMoveValueResponse) -> Result<Self> {
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<ResponseHeader> {
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<KeyValue> {
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());

View File

@@ -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<Vec<String>> {
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<Vec<String>> {
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::<Vec<_>>()
.await?;
.collect::<CatalogResult<Vec<_>>>()?;
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()

View File

@@ -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]

View File

@@ -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)?;

View File

@@ -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

View File

@@ -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]

View File

@@ -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"

View File

@@ -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);

View File

@@ -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<MetaPeerClient>;
@@ -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<Vec<u8>>) -> Result<Vec<KeyValue>> {
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;

View File

@@ -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(),
}
}

View File

@@ -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;

View File

@@ -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<u8> = 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<u8> = 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

View File

@@ -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,

View File

@@ -40,6 +40,4 @@ pub mod table_routes;
#[cfg(test)]
mod test_util;
pub mod util;
pub use crate::error::Result;

View File

@@ -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<u8> = 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());
}
}

View File

@@ -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";

View File

@@ -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.

View File

@@ -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,

View File

@@ -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())
}

View File

@@ -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__/";

View File

@@ -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<api::v1::meta::RangeResponse> {
impl KvBackend for Noop {
fn name(&self) -> &str {
"Noop"
}
async fn range(&self, _: RangeRequest) -> Result<RangeResponse> {
unreachable!()
}
async fn put(
&self,
_: api::v1::meta::PutRequest,
) -> Result<api::v1::meta::PutResponse> {
async fn put(&self, _: PutRequest) -> Result<PutResponse> {
unreachable!()
}
async fn batch_put(
&self,
_: api::v1::meta::BatchPutRequest,
) -> Result<api::v1::meta::BatchPutResponse> {
async fn batch_put(&self, _: BatchPutRequest) -> Result<BatchPutResponse> {
unreachable!()
}
@@ -232,27 +232,25 @@ mod tests {
async fn compare_and_put(
&self,
_: CompareAndPutRequest,
) -> Result<api::v1::meta::CompareAndPutResponse> {
Ok(api::v1::meta::CompareAndPutResponse::default())
) -> Result<CompareAndPutResponse> {
Ok(CompareAndPutResponse::default())
}
async fn delete_range(
&self,
_: api::v1::meta::DeleteRangeRequest,
) -> Result<api::v1::meta::DeleteRangeResponse> {
async fn delete_range(&self, _: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
unreachable!()
}
async fn move_value(
&self,
_: api::v1::meta::MoveValueRequest,
) -> Result<api::v1::meta::MoveValueResponse> {
async fn move_value(&self, _: MoveValueRequest) -> Result<MoveValueResponse> {
unreachable!()
}
async fn batch_delete(&self, _: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
unreachable!()
}
fn as_any(&self) -> &dyn Any {
self
}
}
let kv_store = Arc::new(Noop {});

View File

@@ -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());

View File

@@ -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,

View File

@@ -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<BatchGetRequest>) -> GrpcResult<BatchGetResponse> {
async fn batch_get(&self, req: Request<PbBatchGetRequest>) -> GrpcResult<PbBatchGetResponse> {
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<RangeRequest>) -> GrpcResult<RangeResponse> {
async fn range(&self, req: Request<PbRangeRequest>) -> GrpcResult<PbRangeResponse> {
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))
}
}

View File

@@ -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?;

View File

@@ -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<RangeRequest>) -> GrpcResult<RangeResponse> {
async fn range(&self, req: Request<PbRangeRequest>) -> GrpcResult<PbRangeResponse> {
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<PutRequest>) -> GrpcResult<PutResponse> {
async fn put(&self, req: Request<PbPutRequest>) -> GrpcResult<PbPutResponse> {
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<BatchGetRequest>) -> GrpcResult<BatchGetResponse> {
async fn batch_get(&self, req: Request<PbBatchGetRequest>) -> GrpcResult<PbBatchGetResponse> {
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<BatchPutRequest>) -> GrpcResult<BatchPutResponse> {
async fn batch_put(&self, req: Request<PbBatchPutRequest>) -> GrpcResult<PbBatchPutResponse> {
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<BatchDeleteRequest>,
) -> GrpcResult<BatchDeleteResponse> {
req: Request<PbBatchDeleteRequest>,
) -> GrpcResult<PbBatchDeleteResponse> {
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<CompareAndPutRequest>,
) -> GrpcResult<CompareAndPutResponse> {
req: Request<PbCompareAndPutRequest>,
) -> GrpcResult<PbCompareAndPutResponse> {
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<DeleteRangeRequest>,
) -> GrpcResult<DeleteRangeResponse> {
req: Request<PbDeleteRangeRequest>,
) -> GrpcResult<PbDeleteRangeResponse> {
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<MoveValueRequest>) -> GrpcResult<MoveValueResponse> {
async fn move_value(
&self,
req: Request<PbMoveValueRequest>,
) -> GrpcResult<PbMoveValueResponse> {
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();

View File

@@ -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<dyn CheckLeader>;
@@ -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<u8>) -> 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<RangeResponse> {
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::<Vec<_>>();
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::<Vec<_>>();
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<TxnResponse> {
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::<Vec<_>>();
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::<Vec<_>>();
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);

View File

@@ -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<RangeResponse> {
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<RangeResponse> {
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::<Vec<_>>();
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<PutResponse> {
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<BatchGetResponse> {
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<BatchPutResponse> {
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<BatchDeleteResponse> {
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<CompareAndPutResponse> {
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<DeleteRangeResponse> {
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::<Vec<_>>();
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<MoveValueResponse> {
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<crate::service::store::txn::TxnResponse> {
type Error = Error;
async fn txn(&self, txn: KvTxn) -> Result<KvTxnResponse> {
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<u8>,
options: Option<GetOptions>,
}
@@ -498,7 +395,6 @@ impl TryFrom<RangeRequest> for Get {
fn try_from(req: RangeRequest) -> Result<Self> {
let RangeRequest {
header,
key,
range_end,
limit,
@@ -519,7 +415,6 @@ impl TryFrom<RangeRequest> for Get {
}
Ok(Get {
cluster_id: header.map_or(0, |h| h.cluster_id),
key,
options: Some(options),
})
@@ -527,7 +422,6 @@ impl TryFrom<RangeRequest> for Get {
}
struct Put {
cluster_id: u64,
key: Vec<u8>,
value: Vec<u8>,
options: Option<PutOptions>,
@@ -538,7 +432,6 @@ impl TryFrom<PutRequest> for Put {
fn try_from(req: PutRequest) -> Result<Self> {
let PutRequest {
header,
key,
value,
prev_kv,
@@ -550,7 +443,6 @@ impl TryFrom<PutRequest> 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<PutRequest> for Put {
}
struct BatchGet {
cluster_id: u64,
keys: Vec<Vec<u8>>,
options: Option<GetOptions>,
}
@@ -568,12 +459,11 @@ impl TryFrom<BatchGetRequest> for BatchGet {
type Error = error::Error;
fn try_from(req: BatchGetRequest) -> Result<Self> {
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<BatchGetRequest> for BatchGet {
}
struct BatchPut {
cluster_id: u64,
kvs: Vec<KeyValue>,
options: Option<PutOptions>,
}
@@ -590,11 +479,7 @@ impl TryFrom<BatchPutRequest> for BatchPut {
type Error = error::Error;
fn try_from(req: BatchPutRequest) -> Result<Self> {
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<BatchPutRequest> for BatchPut {
}
Ok(BatchPut {
cluster_id: header.map_or(0, |h| h.cluster_id),
kvs,
options: Some(options),
})
@@ -610,7 +494,6 @@ impl TryFrom<BatchPutRequest> for BatchPut {
}
struct BatchDelete {
cluster_id: u64,
keys: Vec<Vec<u8>>,
options: Option<DeleteOptions>,
}
@@ -619,11 +502,7 @@ impl TryFrom<BatchDeleteRequest> for BatchDelete {
type Error = error::Error;
fn try_from(req: BatchDeleteRequest) -> Result<Self> {
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<BatchDeleteRequest> for BatchDelete {
}
Ok(BatchDelete {
cluster_id: header.map_or(0, |h| h.cluster_id),
keys,
options: Some(options),
})
@@ -639,7 +517,6 @@ impl TryFrom<BatchDeleteRequest> for BatchDelete {
}
struct CompareAndPut {
cluster_id: u64,
key: Vec<u8>,
expect: Vec<u8>,
value: Vec<u8>,
@@ -650,15 +527,9 @@ impl TryFrom<CompareAndPutRequest> for CompareAndPut {
type Error = error::Error;
fn try_from(req: CompareAndPutRequest) -> Result<Self> {
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<CompareAndPutRequest> for CompareAndPut {
}
struct Delete {
cluster_id: u64,
key: Vec<u8>,
options: Option<DeleteOptions>,
}
@@ -678,7 +548,6 @@ impl TryFrom<DeleteRangeRequest> for Delete {
fn try_from(req: DeleteRangeRequest) -> Result<Self> {
let DeleteRangeRequest {
header,
key,
range_end,
prev_kv,
@@ -695,7 +564,6 @@ impl TryFrom<DeleteRangeRequest> for Delete {
}
Ok(Delete {
cluster_id: header.map_or(0, |h| h.cluster_id),
key,
options: Some(options),
})
@@ -703,7 +571,6 @@ impl TryFrom<DeleteRangeRequest> for Delete {
}
struct MoveValue {
cluster_id: u64,
from_key: Vec<u8>,
to_key: Vec<u8>,
delete_options: Option<DeleteOptions>,
@@ -713,14 +580,9 @@ impl TryFrom<MoveValueRequest> for MoveValue {
type Error = error::Error;
fn try_from(req: MoveValueRequest) -> Result<Self> {
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();

View File

@@ -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);

View File

@@ -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<u8>) -> Result<Option<KeyValue>>;
/// Check if a key exists, it does not return the value.
async fn exists(&self, key: Vec<u8>) -> Result<bool>;
/// 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<u8>, prev_kv: bool) -> Result<Option<KeyValue>>;
}
#[async_trait::async_trait]
impl<T> KvStoreExt for T
where
T: KvStore + ?Sized,
{
async fn get(&self, key: Vec<u8>) -> Result<Option<KeyValue>> {
let req = RangeRequest {
key,
..Default::default()
};
let mut kvs = self.range(req).await?.kvs;
Ok(kvs.pop())
}
async fn exists(&self, key: Vec<u8>) -> Result<bool> {
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<u8>, prev_kv: bool) -> Result<Option<KeyValue>> {
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());
}
}

View File

@@ -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<dyn KvStore>;
pub type KvStoreRef = Arc<dyn KvBackend<Error = Error>>;
pub type ResettableKvStoreRef = Arc<dyn ResettableKvStore>;
#[async_trait::async_trait]
pub trait KvStore: TxnService {
async fn range(&self, req: RangeRequest) -> Result<RangeResponse>;
async fn put(&self, req: PutRequest) -> Result<PutResponse>;
async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse>;
async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse>;
async fn batch_delete(&self, req: BatchDeleteRequest) -> Result<BatchDeleteResponse>;
async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result<CompareAndPutResponse>;
async fn delete_range(&self, req: DeleteRangeRequest) -> Result<DeleteRangeResponse>;
async fn move_value(&self, req: MoveValueRequest) -> Result<MoveValueResponse>;
}
pub trait ResettableKvStore: KvStore {
pub trait ResettableKvStore: KvBackend<Error = Error> {
fn reset(&self);
}

View File

@@ -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<BTreeMap<Vec<u8>, Vec<u8>>>,
}
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<Error>;
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<RangeResponse> {
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::<Vec<_>>()
};
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<PutResponse> {
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<BatchGetResponse> {
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<BatchPutResponse> {
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<BatchDeleteResponse> {
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<CompareAndPutResponse> {
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<DeleteRangeResponse> {
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::<Vec<_>>()
};
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<MoveValueResponse> {
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<TxnResponse> {
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<KeyValue> {
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();
}
}

View File

@@ -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<Option<TableGlobalValue>> {
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::<Vec<_>>(),
};
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<TableRouteValue> {
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<Option<(Vec<u8>, Vec<u8>)>> {
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<'_> {

View File

@@ -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<Instance>, key: &str) -> Option<Option<Kv>> {
fn get_table_cache(instance: &Arc<Instance>, key: &str) -> Option<Option<KeyValue>> {
let catalog_manager = instance
.catalog_manager()
.as_any()