Merge pull request #17 from Smityz/support-rawkv

support RawKV

Author: andylokandy

Makefile

@@ -17,12 +17,17 @@ pre-build: directories target/tikv_client_glue.cc include/tikv_client_glue.h
 clean:
 	cargo clean
 
-run-example: target/tikv-example
-	RUST_LOG=debug $(cur_makefile_path)/target/tikv-example
+run-raw-example: target/tikv-raw-example
+	RUST_LOG=debug $(cur_makefile_path)/target/tikv-raw-example
 
-target/tikv-example: target/debug/libtikv_client.a example/main.cpp
-	c++ $(cur_makefile_path)/example/main.cpp -o $(cur_makefile_path)/target/tikv-example -std=c++17 -g -I$(cur_makefile_path)/include -L$(cur_makefile_path)/target/debug -ltikv_client -lpthread -ldl -lssl -lcrypto
+run-txn-example: target/tikv-txn-example 
+	RUST_LOG=debug $(cur_makefile_path)/target/tikv-txn-example
 
+target/tikv-raw-example: target/debug/libtikv_client.a example/raw.cpp
+	c++ $(cur_makefile_path)/example/raw.cpp -o $(cur_makefile_path)/target/tikv-raw-example -std=c++17 -g -I$(cur_makefile_path)/include -L$(cur_makefile_path)/target/debug -ltikv_client -lpthread -ldl -lssl -lcrypto
+
+target/tikv-txn-example: target/debug/libtikv_client.a example/txn.cpp
+	c++ $(cur_makefile_path)/example/txn.cpp -o $(cur_makefile_path)/target/tikv-txn-example -std=c++17 -g -I$(cur_makefile_path)/include -L$(cur_makefile_path)/target/debug -ltikv_client -lpthread -ldl -lssl -lcrypto
 
 target/tikv_client_glue.cc: src/lib.rs
 	cxxbridge $(cur_makefile_path)/src/lib.rs > $(cur_makefile_path)/target/tikv_client_glue.cc

README.md

@@ -28,5 +28,8 @@ make release
 
 ```bash
 tiup playground nightly
-make run-example
+# run rawkv example
+make run-raw-example
+# run txnkv example
+make run-txn-example        
 ```

example/raw.cpp

@@ -0,0 +1,40 @@
+// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
+
+#include "tikv_client.h"
+#include <iostream>
+
+int main() {
+    auto client = tikv_client::RawKVClient({"127.0.0.1:2379"});
+
+    const std::uint32_t kTimeoutMs = 10;
+    client.put("k1", "v1", kTimeoutMs);
+
+    auto val = client.get("k1",kTimeoutMs);
+    if (val) {
+        std::cout << "get key: \n(k1:" << *val << ")" << std::endl;
+    } else {
+        std::cout << "key not found" << std::endl;
+    }
+
+    client.batch_put({{"k2","v2"},{"k3","v3"},{"k4","v4"},{"k5","v5"}}, kTimeoutMs);
+
+    const std::uint32_t kLimit = 20;
+    // scan [k1,k6), limit 20, timeout 10ms
+    auto kv_pairs = client.scan("k1","k6", kLimit ,kTimeoutMs);
+    std::cout<<"scan[\"k1\",\"k6\"):"<<std::endl;
+    for (auto iter = kv_pairs.begin(); iter != kv_pairs.end(); ++iter) {
+        std::cout << "(" << iter->key << ": " << iter->value  << ") ";
+    }
+    std::cout << std::endl;
+
+    // delete [k3,k5), so [k1,k6) should be [k1,k3) + [k5,k6)
+    std::cout<<"scan[\"k1\",\"k6\") after delete:"<<std::endl;
+    client.remove_range("k3","k5",kTimeoutMs);
+    kv_pairs = client.scan("k1","k6", kLimit ,kTimeoutMs);
+    for (auto iter = kv_pairs.begin(); iter != kv_pairs.end(); ++iter) {
+        std::cout << "(" << iter->key << ": " << iter->value  << ") ";
+    }
+    std::cout << std::endl;
+
+    return 0;
+}

example/main.cpp renamed to example/txn.cpp

@@ -14,6 +14,7 @@ struct KvPair final {
     std::string value;
 
     KvPair(std::string &&key, std::string &&value);
+    ffi::KvPair to_ffi();
 };
 
 class Transaction {
@@ -42,6 +43,20 @@ class TransactionClient {
     ::rust::cxxbridge1::Box<tikv_client_glue::TransactionClient> _client;
 };
 
-}
+class RawKVClient {
+public:
+    RawKVClient(const std::vector<std::string> &pd_endpoints);
+    std::optional<std::string> get(const std::string &key,const std::uint32_t timeout);
+    void put(const std::string &key, const std::string &value, const std::uint32_t timeout);
+    void batch_put(const std::vector<KvPair> &kvs, const std::uint32_t timeout);
+    void remove(const std::string &key, const std::uint32_t timeout);
+    void remove_range(const std::string &start_key, const std::string &end_key, const std::uint32_t timeout);
+    std::vector<KvPair> scan(const std::string &startKey, const std::string &endKey, std::uint32_t limit, const std::uint32_t timeout);
+
+private:
+    ::rust::cxxbridge1::Box<tikv_client_glue::RawKVClient> _client;
+};
+
+} // namespace tikv_client
 
 #endif //_TIKV_CLIENT_H_

include/tikv_client.h

@@ -573,12 +573,15 @@ Vec<T>::Vec(unsafe_bitcopy_t, const Vec &bits) noexcept : repr(bits.repr) {}
 } // namespace rust
 
 struct Key;
-struct KvPair;
 struct OptionalValue;
 enum class Bound : ::std::uint8_t;
+namespace ffi {
+  struct KvPair;
+}
 namespace tikv_client_glue {
   struct TransactionClient;
   struct Transaction;
+  struct RawKVClient;
 }
 
 #ifndef CXXBRIDGE1_STRUCT_Key
@@ -590,15 +593,17 @@ struct Key final {
 };
 #endif // CXXBRIDGE1_STRUCT_Key
 
-#ifndef CXXBRIDGE1_STRUCT_KvPair
-#define CXXBRIDGE1_STRUCT_KvPair
+namespace ffi {
+#ifndef CXXBRIDGE1_STRUCT_ffi$KvPair
+#define CXXBRIDGE1_STRUCT_ffi$KvPair
 struct KvPair final {
   ::rust::Vec<::std::uint8_t> key;
   ::rust::Vec<::std::uint8_t> value;
 
   using IsRelocatable = ::std::true_type;
 };
-#endif // CXXBRIDGE1_STRUCT_KvPair
+#endif // CXXBRIDGE1_STRUCT_ffi$KvPair
+} // namespace ffi
 
 #ifndef CXXBRIDGE1_STRUCT_OptionalValue
 #define CXXBRIDGE1_STRUCT_OptionalValue
@@ -620,6 +625,20 @@ enum class Bound : ::std::uint8_t {
 #endif // CXXBRIDGE1_ENUM_Bound
 
 namespace tikv_client_glue {
+::rust::Box<::tikv_client_glue::RawKVClient> raw_client_new(const ::std::vector<::std::string> &pd_endpoints);
+
+::OptionalValue raw_get(const ::tikv_client_glue::RawKVClient &client, const ::std::string &key, ::std::uint32_t timeout_ms);
+
+void raw_put(const ::tikv_client_glue::RawKVClient &cli, const ::std::string &key, const ::std::string &val, ::std::uint32_t timeout_ms);
+
+::rust::Vec<::ffi::KvPair> raw_scan(const ::tikv_client_glue::RawKVClient &cli, const ::std::string &start, const ::std::string &end, ::std::uint32_t limit, ::std::uint32_t timeout_ms);
+
+void raw_delete(const ::tikv_client_glue::RawKVClient &cli, const ::std::string &key, ::std::uint32_t timeout_ms);
+
+void raw_delete_range(const ::tikv_client_glue::RawKVClient &cli, const ::std::string &startKey, const ::std::string &endKey, ::std::uint32_t timeout_ms);
+
+void raw_batch_put(const ::tikv_client_glue::RawKVClient &cli, const ::std::vector<::ffi::KvPair> &pairs, ::std::uint32_t timeout_ms);
+
 ::rust::Box<::tikv_client_glue::TransactionClient> transaction_client_new(const ::std::vector<::std::string> &pd_endpoints);
 
 ::rust::Box<::tikv_client_glue::Transaction> transaction_client_begin(const ::tikv_client_glue::TransactionClient &client);
@@ -630,11 +649,11 @@ ::OptionalValue transaction_get(const ::tikv_client_glue::Transaction &transacti
 
 ::OptionalValue transaction_get_for_update(::tikv_client_glue::Transaction &transaction, const ::std::string &key);
 
-::rust::Vec<::KvPair> transaction_batch_get(::tikv_client_glue::Transaction &transaction, const ::std::vector<::std::string> &keys);
+::rust::Vec<::ffi::KvPair> transaction_batch_get(::tikv_client_glue::Transaction &transaction, const ::std::vector<::std::string> &keys);
 
-::rust::Vec<::KvPair> transaction_batch_get_for_update(::tikv_client_glue::Transaction &transaction, const ::std::vector<::std::string> &keys);
+::rust::Vec<::ffi::KvPair> transaction_batch_get_for_update(::tikv_client_glue::Transaction &transaction, const ::std::vector<::std::string> &keys);
 
-::rust::Vec<::KvPair> transaction_scan(::tikv_client_glue::Transaction &transaction, const ::std::string &start, ::Bound start_bound, const ::std::string &end, ::Bound end_bound, ::std::uint32_t limit);
+::rust::Vec<::ffi::KvPair> transaction_scan(::tikv_client_glue::Transaction &transaction, const ::std::string &start, ::Bound start_bound, const ::std::string &end, ::Bound end_bound, ::std::uint32_t limit);
 
 ::rust::Vec<::Key> transaction_scan_keys(::tikv_client_glue::Transaction &transaction, const ::std::string &start, ::Bound start_bound, const ::std::string &end, ::Bound end_bound, ::std::uint32_t limit);
 

include/tikv_client_glue.h

@@ -15,6 +15,7 @@ mod ffi {
         key: Vec<u8>,
     }
 
+    #[namespace = "ffi"]
     struct KvPair {
         key: Vec<u8>,
         value: Vec<u8>,
@@ -35,6 +36,42 @@ mod ffi {
     extern "Rust" {
         type TransactionClient;
         type Transaction;
+        type RawKVClient;
+
+        fn raw_client_new(pd_endpoints: &CxxVector<CxxString>) -> Result<Box<RawKVClient>>;
+
+        fn raw_get(client: &RawKVClient, key: &CxxString, timeout_ms: u32)
+            -> Result<OptionalValue>;
+
+        fn raw_put(
+            cli: &RawKVClient,
+            key: &CxxString,
+            val: &CxxString,
+            timeout_ms: u32,
+        ) -> Result<()>;
+
+        fn raw_scan(
+            cli: &RawKVClient,
+            start: &CxxString,
+            end: &CxxString,
+            limit: u32,
+            timeout_ms: u32,
+        ) -> Result<Vec<KvPair>>;
+
+        fn raw_delete(cli: &RawKVClient, key: &CxxString, timeout_ms: u32) -> Result<()>;
+
+        fn raw_delete_range(
+            cli: &RawKVClient,
+            startKey: &CxxString,
+            endKey: &CxxString,
+            timeout_ms: u32,
+        ) -> Result<()>;
+
+        fn raw_batch_put(
+            cli: &RawKVClient,
+            pairs: &CxxVector<KvPair>,
+            timeout_ms: u32,
+        ) -> Result<()>;
 
         fn transaction_client_new(
             pd_endpoints: &CxxVector<CxxString>,
@@ -98,11 +135,28 @@ struct TransactionClient {
     inner: tikv_client::TransactionClient,
 }
 
+struct RawKVClient {
+    inner: tikv_client::RawClient,
+}
+
 #[repr(transparent)]
 struct Transaction {
     inner: tikv_client::Transaction,
 }
 
+fn raw_client_new(pd_endpoints: &CxxVector<CxxString>) -> Result<Box<RawKVClient>> {
+    env_logger::init();
+
+    let pd_endpoints = pd_endpoints
+        .iter()
+        .map(|str| str.to_str().map(ToOwned::to_owned))
+        .collect::<std::result::Result<Vec<_>, _>>()?;
+
+    Ok(Box::new(RawKVClient {
+        inner: block_on(tikv_client::RawClient::new(pd_endpoints))?,
+    }))
+}
+
 fn transaction_client_new(pd_endpoints: &CxxVector<CxxString>) -> Result<Box<TransactionClient>> {
     env_logger::init();
 
@@ -128,6 +182,73 @@ fn transaction_client_begin_pessimistic(client: &TransactionClient) -> Result<Bo
     }))
 }
 
+fn raw_get(cli: &RawKVClient, key: &CxxString, timeout_ms: u32) -> Result<OptionalValue> {
+    match block_on(cli.inner.get(key.as_bytes().to_vec()))? {
+        Some(value) => Ok(OptionalValue {
+            is_none: false,
+            value,
+        }),
+        None => Ok(OptionalValue {
+            is_none: true,
+            value: Vec::new(),
+        }),
+    }
+}
+
+// TODO(smityz): implement timeout
+fn raw_put(cli: &RawKVClient, key: &CxxString, val: &CxxString, timeout_ms: u32) -> Result<()> {
+    block_on(
+        cli.inner
+            .put(key.as_bytes().to_vec(), val.as_bytes().to_vec()),
+    )?;
+    Ok(())
+}
+
+fn raw_scan(
+    cli: &RawKVClient,
+    start: &CxxString,
+    end: &CxxString,
+    limit: u32,
+    timeout_ms: u32,
+) -> Result<Vec<KvPair>> {
+    let rg = to_bound_range(start, Bound::Included, end, Bound::Excluded);
+    let pairs = block_on(cli.inner.scan(rg, limit))?
+        .into_iter()
+        .map(|tikv_client::KvPair(key, value)| KvPair {
+            key: key.into(),
+            value,
+        })
+        .collect();
+    Ok(pairs)
+}
+
+fn raw_delete(cli: &RawKVClient, key: &CxxString, timeout_ms: u32) -> Result<()> {
+    block_on(cli.inner.delete(key.as_bytes().to_vec()))?;
+    Ok(())
+}
+
+fn raw_delete_range(
+    cli: &RawKVClient,
+    start_key: &CxxString,
+    end_key: &CxxString,
+    timeout_ms: u32,
+) -> Result<()> {
+    let rg = to_bound_range(start_key, Bound::Included, end_key, Bound::Excluded);
+    block_on(cli.inner.delete_range(rg))?;
+    Ok(())
+}
+
+fn raw_batch_put(cli: &RawKVClient, pairs: &CxxVector<KvPair>, timeout_ms: u32) -> Result<()> {
+    let tikv_pairs: Vec<tikv_client::KvPair> = pairs
+        .iter()
+        .map(|KvPair { key, value }| -> tikv_client::KvPair {
+            tikv_client::KvPair(key.to_vec().into(), value.to_vec())
+        })
+        .collect();
+    block_on(cli.inner.batch_put(tikv_pairs))?;
+    Ok(())
+}
+
 fn transaction_get(transaction: &Transaction, key: &CxxString) -> Result<OptionalValue> {
     match block_on(transaction.inner.get(key.as_bytes().to_vec()))? {
         Some(value) => Ok(OptionalValue {

src/lib.rs

@@ -12,9 +12,68 @@ KvPair::KvPair(std::string &&key, std::string &&value)
     , value(std::move(value))
 {}
 
+ffi::KvPair KvPair::to_ffi() {
+    ffi::KvPair f_pair;
+    f_pair.key.reserve(key.size());
+    for (const auto &c : this->key) {
+        f_pair.key.emplace_back(static_cast<std::uint8_t>(c));
+    }
+    f_pair.value.reserve(value.size());
+    for (const auto &c : this->value) {
+        f_pair.value.emplace_back(static_cast<std::uint8_t>(c));
+    }
+    return f_pair;
+}
+
 TransactionClient::TransactionClient(const std::vector<std::string> &pd_endpoints):
     _client(tikv_client_glue::transaction_client_new(pd_endpoints)) {}
 
+RawKVClient::RawKVClient(const std::vector<std::string> &pd_endpoints):
+    _client(tikv_client_glue::raw_client_new(pd_endpoints)) {}
+
+std::optional<std::string> RawKVClient::get(const std::string &key, const std::uint32_t timeout) {
+    auto val = tikv_client_glue::raw_get(*_client,key,timeout);
+    if (val.is_none) {
+        return std::nullopt;
+    } else {
+        return std::string{val.value.begin(), val.value.end()};
+    }
+}
+
+void RawKVClient::put(const std::string &key, const std::string &value, const std::uint32_t timeout) {
+    tikv_client_glue::raw_put(*_client,key,value,timeout);
+}
+
+void RawKVClient::batch_put(const std::vector<KvPair> &kv_pairs, const std::uint32_t timeout) {
+    std::vector<ffi::KvPair> pairs;
+    pairs.reserve(kv_pairs.size());
+    for (auto pair: kv_pairs) {
+        pairs.emplace_back(pair.to_ffi());
+    }
+    tikv_client_glue::raw_batch_put(*_client,pairs,timeout);
+}
+
+std::vector<KvPair> RawKVClient::scan(const std::string &startKey, const std::string &endKey, std::uint32_t limit, const std::uint32_t timeout){
+    auto kv_pairs = tikv_client_glue::raw_scan(*_client,startKey,endKey,limit,timeout);  
+    std::vector<KvPair> result;
+    result.reserve(kv_pairs.size());
+    for (auto iter = kv_pairs.begin(); iter != kv_pairs.end(); ++iter) {
+        result.emplace_back(
+            std::string{(iter->key).begin(), (iter->key).end()},
+            std::string{(iter->value).begin(), (iter->value).end()}
+        );
+    }
+    return result;
+}
+
+void RawKVClient::remove(const std::string &key, const std::uint32_t timeout) {
+    tikv_client_glue::raw_delete(*_client,key,timeout);
+}
+
+void RawKVClient::remove_range(const std::string &start_key,const std::string &end_key, const std::uint32_t timeout) {
+    tikv_client_glue::raw_delete_range(*_client,start_key,end_key,timeout);
+}
+
 Transaction TransactionClient::begin() {
     return Transaction(transaction_client_begin(*_client));
 }