From 29d9d90c8b2bd1ee205f41bdf7ba6d0b601c6651 Mon Sep 17 00:00:00 2001 From: lance6716 Date: Wed, 14 Oct 2026 17:13:28 +0800 Subject: [PATCH] feat: add async C and C++ bindings for TiKV RawKV Signed-off-by: lance6716 --- Cargo.toml | 1 + bindings/c/Cargo.toml | 15 + bindings/c/README.md | 390 +++ bindings/c/examples/metrics_snapshot_demo.cpp | 189 ++ bindings/c/examples/raw_async_demo.c | 322 +++ bindings/c/examples/raw_async_demo.cpp | 449 +++ bindings/c/include/tikv_client_async.hpp | 1854 ++++++++++++ bindings/c/include/tikv_client_async_c.h | 502 ++++ bindings/c/include/tikv_client_metrics.hpp | 620 +++++ bindings/c/include/tikv_client_metrics_c.h | 116 + bindings/c/src/lib.rs | 2475 +++++++++++++++++ 11 files changed, 6933 insertions(+) create mode 100644 bindings/c/Cargo.toml create mode 100644 bindings/c/README.md create mode 100644 bindings/c/examples/metrics_snapshot_demo.cpp create mode 100644 bindings/c/examples/raw_async_demo.c create mode 100644 bindings/c/examples/raw_async_demo.cpp create mode 100644 bindings/c/include/tikv_client_async.hpp create mode 100644 bindings/c/include/tikv_client_async_c.h create mode 100644 bindings/c/include/tikv_client_metrics.hpp create mode 100644 bindings/c/include/tikv_client_metrics_c.h create mode 100644 bindings/c/src/lib.rs diff --git a/Cargo.toml b/Cargo.toml index ca1029f3..5565387f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ name = "tikv_client" [workspace] members = [ ".", + "bindings/c", "proto-build", ] diff --git a/bindings/c/Cargo.toml b/bindings/c/Cargo.toml new file mode 100644 index 00000000..a874a582 --- /dev/null +++ b/bindings/c/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "tikv-client-rust-c" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +name = "tikv_client_rust_c" +crate-type = ["cdylib", "staticlib"] + +[dependencies] +futures = "0.3" +prometheus = { version = "0.13", default-features = false } +tikv-client = { path = "../..", default-features = false } +tokio = { version = "1", features = ["rt-multi-thread", "net", "time"] } diff --git a/bindings/c/README.md b/bindings/c/README.md new file mode 100644 index 00000000..a93d038a --- /dev/null +++ b/bindings/c/README.md @@ -0,0 +1,390 @@ +# TiKV Rust Client C Binding + +这个目录提供 TiKV Rust Client 的异步 C ABI 和 C++ 头文件封装。 + +接口使用异步请求路径： + +- 调用方提交请求时只做参数校验和必要 copy +- 真正的 TiKV 请求在 Rust async runtime 里执行 +- 完成信号通过 completion callback + userdata 收口，而不是 `block_on` +- shutdown / drain / cancel 语义在 FFI 边界显式定义 + +## 当前覆盖能力 + +- `connect` +- `put` +- `put_with_ttl` +- `get` +- `get_key_ttl` +- `delete` +- `delete_range` +- `scan` +- `scan_reverse` +- `batch_write`（可混合 `put/delete`） +- `batch_write_with_ttl`（每个 `put` 独立设置 TTL） +- `batch_get` +- `shutdown_begin` +- `drain` +- `request_cancel` + +## 构建 + +```bash +cargo build -p tikv-client-rust-c +``` + +输出产物位于： + +- `target/debug/libtikv_client_rust_c.dylib` 或 Linux 下的 `libtikv_client_rust_c.so` +- `target/debug/libtikv_client_rust_c.a` +- 异步客户端头文件：`bindings/c/include/tikv_client_async_c.h` +- 指标快照头文件：`bindings/c/include/tikv_client_metrics_c.h` +- C++ 指标封装头文件：`bindings/c/include/tikv_client_metrics.hpp` + +## 编译 C 示例 + +macOS: + +```bash +clang \ + -Ibindings/c/include \ + bindings/c/examples/raw_async_demo.c \ + -Ltarget/debug \ + -ltikv_client_rust_c \ + -Wl,-rpath,$PWD/target/debug \ + -o target/debug/raw_async_demo +``` + +Linux: + +```bash +cc \ + -Ibindings/c/include \ + bindings/c/examples/raw_async_demo.c \ + -Ltarget/debug \ + -ltikv_client_rust_c \ + -Wl,-rpath,'$ORIGIN' \ + -o target/debug/raw_async_demo +``` +## 接口语义 + +- Rust 客户端的 gRPC response decode 上限仍默认为 4 MiB。为保持原有 C/C++ + 行为，`tikv_async_client_new` 和 C++ `Client::connect` / `connect_as_shared` + 默认使用 `SIZE_MAX`（Rust 的 `usize::MAX`）。如需设置上限，在创建客户端时 + 显式指定字节数；显式传入 `0` 则使用 Rust 的 4 MiB 默认值。C 接口可调用 + `tikv_async_client_new_with_grpc_max_decoding_message_size(pd, timeout, keyspace, 32 * 1024 * 1024, &client, &error)`； + C++ 接口可调用 `Client::connect(pd, timeout, keyspace, 32 * 1024 * 1024)`。 +- submit 返回 `TIKV_STATUS_OK` 只表示请求已被接收，且之后一定 callback 一次 +- submit 返回非 `TIKV_STATUS_OK` 表示请求未被接收，不会触发 callback +- completion callback 运行在 Rust runtime worker 线程，不保证回到 submit 线程 +- `tikv_async_request_destroy` 可以早于 callback 调用；Rust 会把内部状态保活到 callback 返回 +- `tikv_async_request_cancel` 是 best-effort 逻辑取消；一旦请求最终被标记为取消，它仍然会通过 callback 以 `TIKV_STATUS_CANCELLED` 显式完成 +- `tikv_async_client_shutdown_begin` 立即拒绝新 submit；若 `cancel_inflight=true`，则所有 in-flight 请求最终通过 callback 以 `TIKV_STATUS_CANCELLED` 收口 +- `tikv_async_client_drain` 等待所有已接收请求完成 callback +- `tikv_async_client_destroy` 会执行 shutdown-begin(cancel) + drain，再释放 client +- `batch_write` 允许一个 batch 内混合 `put/delete`；底层仍复用 `client-rust` 的 `batch_put` / `batch_delete`，但保持调用方的 put/delete phase 顺序 +- 连续 `put` 段内若同一 key 重复出现，binding 会在该位置切段，避免同一 key 的多次写入被压进同一个底层 `batch_put` +- `batch_get` 只返回存在的 key，且结果顺序不保证与输入一致 +- `scan_reverse` 映射到底层 `scan_reverse`；由于 `client-rust` 现有限制，空 `start_key` 不被支持 + +## 编译 C++ 示例 + +macOS: + +```bash +clang++ \ + -std=c++17 \ + -Ibindings/c/include \ + bindings/c/examples/raw_async_demo.cpp \ + -Ltarget/debug \ + -ltikv_client_rust_c \ + -Wl,-rpath,$PWD/target/debug \ + -o target/debug/raw_async_demo_cpp +``` + +Linux: + +```bash +c++ \ + -std=c++17 \ + -Ibindings/c/include \ + bindings/c/examples/raw_async_demo.cpp \ + -Ltarget/debug \ + -ltikv_client_rust_c \ + -Wl,-rpath,'$ORIGIN' \ + -o target/debug/raw_async_demo_cpp +``` + +## 编译 Metrics C++ 示例 + +macOS: + +```bash +clang++ \ + -std=c++17 \ + -Ibindings/c/include \ + bindings/c/examples/metrics_snapshot_demo.cpp \ + -Ltarget/debug \ + -ltikv_client_rust_c \ + -Wl,-rpath,$PWD/target/debug \ + -o target/debug/metrics_snapshot_demo +``` + +Linux: + +```bash +c++ \ + -std=c++17 \ + -Ibindings/c/include \ + bindings/c/examples/metrics_snapshot_demo.cpp \ + -Ltarget/debug \ + -ltikv_client_rust_c \ + -Wl,-rpath,'$ORIGIN' \ + -o target/debug/metrics_snapshot_demo +``` + +Metrics C++ 封装中，`MetricsSnapshot::get()` 返回 owning +`MetricSample`，默认迭代器产出 `Result`，可以安全传给其它 +模块；只在当前 snapshot 生命周期内临时查看时，可以显式使用 +`MetricsSnapshot::get_view()` 或 `MetricsSnapshot::views()` 获取零拷贝的 +`MetricSampleView`。 + +```cpp +for (auto sample_result : snapshot) { + MetricSample sample = sample_result.take(); +} + +for (auto view_result : snapshot.views()) { + MetricSampleView view = view_result.value(); +} +``` + +## C++ 使用方法 + + +### 创建 Client + +```cpp +Result client_result = Client::connect("127.0.0.1:2379", 5); + +if (!client_result.ok()) { + std::cerr << client_result.status().to_string() << '\n'; + return 1; +} + +Client client = client_result.take(); +``` + +多个对象共享同一个 client： + +```cpp +Result client_result = + Client::connect_as_shared("127.0.0.1:2379", 5); + +if (!client_result.ok()) { + std::cerr << client_result.status().to_string() << '\n'; + return 1; +} + +Client::SharedPtr client = client_result.value(); +``` + +### 写入和读取 + +```cpp +// 写入 +AsyncRequest put_request = + client.put("cpp-demo/key", "cpp-demo/value"); + +Status put_status = put_request.wait(); +if (!put_status.ok()) { + std::cerr << put_status.to_string() << '\n'; + return 1; +} +``` + +```cpp +//读取 +AsyncRequest get_request = client.get("cpp-demo/key"); + +Result value_result = get_request.wait(); +if (!value_result.ok()) { + std::cerr << value_result.status().to_string() << '\n'; + return 1; +} + +const OptionalValue &value = value_result.value(); +if (value.has_value()) { + std::cout << value->to_string() << '\n'; +} +``` + +### 自定义 callback + +```cpp +struct CallbackState { + std::promise done; +}; + +void on_complete(const Status &status, void *userdata) noexcept { + CallbackState *state = static_cast(userdata); + + // callback 运行在 Rust runtime worker 线程上，不能抛异常。 + // 如果要把状态传回调用线程，建议用 promise、队列或其他线程安全机制。 + try { + state->done.set_value(status); + } catch (...) { + } +} + +CallbackState callback_state; +std::future callback_done = callback_state.done.get_future(); + +SubmitOptions options; +options.callback = on_complete; +options.userdata = &callback_state; + +// userdata 不会被 C++ wrapper 接管，调用方必须保证它活到 callback 结束。 +AsyncRequest put_request = + client.put("cpp-demo/callback-key", "cpp-demo/callback-value", options); + +if (!put_request.submitted()) { + std::cerr << put_request.submit_status().to_string() << '\n'; + return 1; +} + +// 对于已提交成功的请求，callback 会被调用一次。 +// 这里等待 callback 自己发出的通知，而不是只等待 AsyncRequest::wait()。 +Status callback_status = callback_done.get(); +if (!callback_status.ok()) { + std::cerr << callback_status.to_string() << '\n'; + return 1; +} +``` + +带返回值的请求也可以使用 callback；callback 只负责通知终态，结果仍然从 +`AsyncRequest` 取： + +```cpp +CallbackState callback_state; +std::future callback_done = callback_state.done.get_future(); + +SubmitOptions options; +options.callback = on_complete; +options.userdata = &callback_state; + +AsyncRequest get_request = + client.get("cpp-demo/key", options); + +if (!get_request.submitted()) { + std::cerr << get_request.submit_status().to_string() << '\n'; + return 1; +} + +Status callback_status = callback_done.get(); +if (!callback_status.ok()) { + std::cerr << callback_status.to_string() << '\n'; + return 1; +} + +Result value_result = get_request.take(); +if (!value_result.ok()) { + std::cerr << value_result.status().to_string() << '\n'; + return 1; +} + +const OptionalValue &value = value_result.value(); +if (value.has_value()) { + std::cout << value->to_string() << '\n'; +} +``` + +### 批量读取 + +```cpp +std::vector keys = { + Slice("cpp-demo/key"), + Slice("cpp-demo/missing"), +}; + +AsyncRequest batch_get_request = client.batch_get(keys); + +Result batch_get_result = batch_get_request.wait(); +if (!batch_get_result.ok()) { + std::cerr << batch_get_result.status().to_string() << '\n'; + return 1; +} + +KvPairs pairs = batch_get_result.take(); +for (std::size_t index = 0; index < pairs.size(); ++index) { + KvPairView pair = pairs[index]; + std::cout << pair.key.to_string() << " = " + << pair.value.to_string() << '\n'; +} +``` + +### 扫描 + +```cpp +AsyncRequest scan_request = + client.scan("cpp-demo/", "cpp-demo0", 128); + +Result scan_result = scan_request.wait(); +if (!scan_result.ok()) { + std::cerr << scan_result.status().to_string() << '\n'; + return 1; +} + +KvPairs pairs = scan_result.take(); +for (std::size_t index = 0; index < pairs.size(); ++index) { + KvPairView pair = pairs[index]; + std::cout << pair.key.to_string() << " = " + << pair.value.to_string() << '\n'; +} +``` + +### 关闭 Client + +```cpp +Status shutdown_status = client.shutdown_begin(false); +if (!shutdown_status.ok()) { + std::cerr << shutdown_status.to_string() << '\n'; + return 1; +} + +Status drain_status = client.drain(); +if (!drain_status.ok()) { + std::cerr << drain_status.to_string() << '\n'; + return 1; +} +``` + +## 运行 + +先准备一个可访问的 PD endpoint，例如 `127.0.0.1:2379`： + +```bash +./target/debug/raw_async_demo 127.0.0.1:2379 +./target/debug/raw_async_demo_cpp 127.0.0.1:2379 +./target/debug/metrics_snapshot_demo 127.0.0.1:2379 +``` + +C++ demo 的参数依次为 PD endpoint、是否使用 default keyspace，以及是否运行 +native TTL 集成测试。前两个参数的既有行为保持不变；第三个参数默认为 +`false`。开启 TTL 集成测试时，TiKV 必须启用 RawKV TTL，例如使用仓库中的 +`config/tikv.toml`（包含 `storage.enable-ttl = true`）。API v2 集群还应使用 +default keyspace： + +```bash +./target/debug/raw_async_demo_cpp 127.0.0.1:2379 true true +``` + +TTL 集成测试覆盖单条 TTL 写入与查询、用 `ttl_secs = 0` 清除 TTL、带不同 +TTL 的批量 put/delete、真实到期，以及无 TTL key 在等待期间持续存在。测试 +使用最长 10 秒的轮询等待到期，不依赖固定时长的 sleep。 + +`metrics_snapshot_demo` 会在同一进程内发起一次 raw put/get/delete， +然后通过 `tikv_metrics_snapshot_collect` 读取当前进程的 Prometheus +registry，并通过 C++ `MetricName` / `MetricTypeLabel` 封装验证 +`tikv_request_total{type="raw_get"}` 以及 raw_get duration histogram +样本是否存在。 diff --git a/bindings/c/examples/metrics_snapshot_demo.cpp b/bindings/c/examples/metrics_snapshot_demo.cpp new file mode 100644 index 00000000..7a203635 --- /dev/null +++ b/bindings/c/examples/metrics_snapshot_demo.cpp @@ -0,0 +1,189 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#include "tikv_client_async.hpp" +#include "tikv_client_metrics.hpp" + +#include +#include +#include + +namespace { + +using tikv::async::AsyncRequest; +using tikv::async::Client; +using tikv::async::OptionalValue; +using tikv::async::Result; +using tikv::async::Slice; +using tikv::async::Status; +using tikv::async::Unit; +using tikv::metrics::MetricSample; +using tikv::metrics::MetricsSnapshot; +namespace metric_names = tikv::metrics::metric_names; +namespace metric_type_labels = tikv::metrics::metric_type_labels; + +bool expect_ok(const char *label, const Status &status) { + if (status.ok()) { + return true; + } + + std::cerr << label << ": " << status.to_string() << '\n'; + return false; +} + +template +bool expect_result(const char *label, const Result &result) { + if (result.ok()) { + return true; + } + + std::cerr << label << ": " << result.status().to_string() << '\n'; + return false; +} + +template +bool expect_submitted(const char *label, const AsyncRequest &request) { + if (request.submitted()) { + return true; + } + + std::cerr << label << ": " << request.submit_status().to_string() << '\n'; + return false; +} + +bool wait_unit(const char *submit_label, + const char *wait_label, + AsyncRequest &request) { + return expect_submitted(submit_label, request) && + expect_ok(wait_label, request.wait()); +} + +bool parse_bool_arg(const char *value) { + return std::strcmp(value, "1") == 0 || std::strcmp(value, "true") == 0 || + std::strcmp(value, "TRUE") == 0 || + std::strcmp(value, "default-keyspace") == 0; +} + +bool run_raw_get(Client *client) { + const Slice key("ffi-metrics-demo/key"); + const Slice value("ffi-metrics-demo/value"); + + auto cleanup_before = client->remove(key); + if (!wait_unit("raw_delete_submit before put", + "wait raw_delete before put", + cleanup_before)) { + return false; + } + + auto put = client->put(key, value); + if (!wait_unit("raw_put_submit", "wait raw_put", put)) { + return false; + } + + auto get = client->get(key); + if (!expect_submitted("raw_get_submit", get)) { + return false; + } + + auto get_result = get.wait(); + if (!expect_result("wait raw_get", get_result)) { + return false; + } + + OptionalValue found = get_result.take(); + if (!found.has_value()) { + std::cerr << "raw_get returned no value\n"; + return false; + } + + std::cout << "raw_get value=" << found->to_string() << '\n'; + + auto cleanup_after = client->remove(key); + return wait_unit("raw_delete_submit after get", + "wait raw_delete after get", + cleanup_after); +} + +} // namespace + +int main(int argc, char **argv) { + const char *pd = argc > 1 ? argv[1] : "127.0.0.1:2379"; + const bool use_default_keyspace = argc > 2 && parse_bool_arg(argv[2]); + + auto client_result = Client::connect(pd, 5, use_default_keyspace); + if (!expect_result("tikv_async_client_new", client_result)) { + return 1; + } + + Client client = client_result.take(); + if (!run_raw_get(&client)) { + return 1; + } + + auto snapshot_result = MetricsSnapshot::collect(); + if (!expect_result("tikv_metrics_snapshot_collect", snapshot_result)) { + return 1; + } + MetricsSnapshot snapshot = snapshot_result.take(); + + bool found_raw_get_total = false; + bool found_raw_get_duration_bucket = false; + bool found_raw_get_duration_count = false; + + for (auto sample_result : snapshot) { + if (!expect_result("tikv_metrics_snapshot_get", sample_result)) { + return 1; + } + MetricSample sample = sample_result.take(); + + if (sample.is_counter(metric_names::TikvRequestTotal, + metric_type_labels::RawGet)) { + found_raw_get_total = true; + std::cout << metric_names::TikvRequestTotal.value() << "{type=\"" + << metric_type_labels::RawGet.value() << "\"} " + << sample.value << '\n'; + } + + if (sample.is_histogram_bucket( + metric_names::TikvRequestDurationSeconds, + metric_type_labels::RawGet)) { + found_raw_get_duration_bucket = true; + std::cout << metric_names::TikvRequestDurationSeconds.value() + << "_bucket{type=\"" << metric_type_labels::RawGet.value() + << "\",le=\"" << sample.le << "\"} " << sample.value + << '\n'; + } + + if (sample.is_histogram_count( + metric_names::TikvRequestDurationSeconds, + metric_type_labels::RawGet)) { + found_raw_get_duration_count = true; + std::cout << metric_names::TikvRequestDurationSeconds.value() + << "_count{type=\"" << metric_type_labels::RawGet.value() + << "\"} " << sample.value << '\n'; + } + } + + if (!found_raw_get_total) { + std::cerr << "No raw_get request counter sample was present.\n"; + } + if (!found_raw_get_duration_bucket) { + std::cerr << "No raw_get request duration bucket was present.\n"; + } + if (!found_raw_get_duration_count) { + std::cerr << "No raw_get request duration count was present.\n"; + } + + Status status = client.shutdown_begin(false); + if (!expect_ok("shutdown_begin", status)) { + return 1; + } + status = client.drain(); + if (!expect_ok("drain", status)) { + return 1; + } + + return found_raw_get_total && found_raw_get_duration_bucket && + found_raw_get_duration_count + ? 0 + : 1; +} diff --git a/bindings/c/examples/raw_async_demo.c b/bindings/c/examples/raw_async_demo.c new file mode 100644 index 00000000..41924b5e --- /dev/null +++ b/bindings/c/examples/raw_async_demo.c @@ -0,0 +1,322 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#include "tikv_client_async_c.h" + +#include +#include +#include +#include + +typedef struct request_latch_t { + volatile int done; + tikv_status_t status; +} request_latch_t; + +static tikv_slice_t slice_from_cstr(const char *value) { + tikv_slice_t slice; + slice.ptr = (const uint8_t *)value; + slice.len = strlen(value); + return slice; +} + +static void print_error_and_destroy(const char *label, tikv_error_t *error) { + const char *message = tikv_error_message(error); + fprintf(stderr, "%s: %s\n", label, message == NULL ? "(null)" : message); + tikv_error_destroy(error); +} + +static int expect_ok(const char *label, + tikv_status_t status, + tikv_error_t *error) { + if (status == TIKV_STATUS_OK) { + return 1; + } + + if (error != NULL) { + print_error_and_destroy(label, error); + } else { + fprintf(stderr, "%s: status=%d\n", label, (int)status); + } + return 0; +} + +static void print_bytes(const uint8_t *ptr, size_t len) { + for (size_t i = 0; i < len; ++i) { + putchar((int)ptr[i]); + } +} + +static void completion_cb(tikv_async_request_t *request, + tikv_status_t final_status, + void *userdata) { + request_latch_t *latch = (request_latch_t *)userdata; + latch->status = final_status; + latch->done = 1; + (void)request; +} + +static int wait_latch_ok(const char *label, + tikv_async_request_t *request, + request_latch_t *latch) { + while (!latch->done) { + struct timespec delay; + delay.tv_sec = 0; + delay.tv_nsec = 10 * 1000 * 1000; + nanosleep(&delay, NULL); + } + + if (latch->status == TIKV_STATUS_OK) { + return 1; + } + + const char *message = tikv_async_request_error_message(request); + fprintf(stderr, + "%s: completion status=%d message=%s\n", + label, + (int)latch->status, + message == NULL ? "(null)" : message); + return 0; +} + +int main(int argc, char **argv) { + const char *pd = argc > 1 ? argv[1] : "127.0.0.1:2379"; + const tikv_slice_t scan_start = slice_from_cstr("tikv-client-rust-c/"); + const tikv_slice_t scan_end = slice_from_cstr("tikv-client-rust-c0"); + const tikv_slice_t key_a = slice_from_cstr("tikv-client-rust-c/a"); + const tikv_slice_t key_b = slice_from_cstr("tikv-client-rust-c/b"); + const tikv_slice_t value_a_v1 = slice_from_cstr("value-a-v1"); + const tikv_slice_t value_a_v2 = slice_from_cstr("value-a-v2"); + const tikv_slice_t value_b = slice_from_cstr("value-b"); + + const tikv_async_raw_mutation_with_ttl_t write_items_buf[] = { + {.op = TIKV_ASYNC_RAW_MUTATION_OP_PUT, + .key = key_a, + .value = value_a_v1, + .ttl_secs = 0}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_PUT, + .key = key_b, + .value = value_b, + .ttl_secs = 0}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + .key = key_a, + .value = {NULL, 0}, + .ttl_secs = 0}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_PUT, + .key = key_a, + .value = value_a_v2, + .ttl_secs = 0}, + }; + const tikv_async_raw_mutations_with_ttl_t write_items = { + .ptr = write_items_buf, + .len = sizeof(write_items_buf) / sizeof(write_items_buf[0]), + }; + const tikv_async_raw_mutation_t cleanup_items_buf[] = { + {.op = TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + .key = key_a, + .value = {NULL, 0}}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + .key = key_b, + .value = {NULL, 0}}, + }; + const tikv_async_raw_mutations_t cleanup_items = { + .ptr = cleanup_items_buf, + .len = sizeof(cleanup_items_buf) / sizeof(cleanup_items_buf[0]), + }; + const tikv_slice_t key_items_buf[] = {key_a, key_b}; + const tikv_slices_t key_items = { + .ptr = key_items_buf, + .len = sizeof(key_items_buf) / sizeof(key_items_buf[0]), + }; + + tikv_async_client_t *client = NULL; + tikv_error_t *error = NULL; + tikv_status_t status = tikv_async_client_new(pd, 5, false, &client, &error); + if (!expect_ok("tikv_async_client_new", status, error)) { + return 1; + } + + request_latch_t reset_latch = {0}; + request_latch_t batch_write_latch = {0}; + request_latch_t batch_get_latch = {0}; + request_latch_t scan_latch = {0}; + request_latch_t reverse_scan_latch = {0}; + request_latch_t cleanup_latch = {0}; + request_latch_t final_scan_latch = {0}; + + tikv_async_request_t *reset = NULL; + tikv_async_request_t *batch_write = NULL; + tikv_async_request_t *batch_get = NULL; + tikv_async_request_t *scan = NULL; + tikv_async_request_t *reverse_scan = NULL; + tikv_async_request_t *cleanup = NULL; + tikv_async_request_t *final_scan = NULL; + + status = tikv_async_raw_delete_range_submit( + client, scan_start, scan_end, completion_cb, &reset_latch, &reset, &error); + if (!expect_ok("delete_range_submit before batch_write", status, error) || + !wait_latch_ok("wait delete_range_submit before batch_write", reset, + &reset_latch)) { + tikv_async_request_destroy(reset); + tikv_async_client_destroy(client); + return 1; + } + tikv_async_request_destroy(reset); + + status = tikv_async_raw_batch_write_with_ttl_submit( + client, write_items, completion_cb, &batch_write_latch, &batch_write, &error); + if (!expect_ok("batch_write_submit", status, error) || + !wait_latch_ok("wait batch_write_submit", batch_write, &batch_write_latch)) { + tikv_async_request_destroy(batch_write); + tikv_async_client_destroy(client); + return 1; + } + tikv_async_request_destroy(batch_write); + + status = tikv_async_raw_batch_get_submit( + client, key_items, completion_cb, &batch_get_latch, &batch_get, &error); + if (!expect_ok("batch_get_submit", status, error)) { + tikv_async_client_destroy(client); + return 1; + } + status = tikv_async_raw_scan_submit( + client, scan_start, scan_end, 16, completion_cb, &scan_latch, &scan, &error); + if (!expect_ok("scan_submit", status, error)) { + tikv_async_request_destroy(batch_get); + tikv_async_client_destroy(client); + return 1; + } + status = tikv_async_raw_scan_reverse_submit(client, + scan_start, + scan_end, + 16, + completion_cb, + &reverse_scan_latch, + &reverse_scan, + &error); + if (!expect_ok("scan_reverse_submit", status, error)) { + tikv_async_request_destroy(batch_get); + tikv_async_request_destroy(scan); + tikv_async_client_destroy(client); + return 1; + } + + if (!wait_latch_ok("wait batch_get_submit", batch_get, &batch_get_latch) || + !wait_latch_ok("wait scan_submit", scan, &scan_latch) || + !wait_latch_ok("wait scan_reverse_submit", reverse_scan, + &reverse_scan_latch)) { + tikv_async_request_destroy(batch_get); + tikv_async_request_destroy(scan); + tikv_async_request_destroy(reverse_scan); + tikv_async_client_destroy(client); + return 1; + } + + tikv_kv_pairs_t pairs = {0}; + status = tikv_async_request_take_batch_get_result(batch_get, &pairs, &error); + if (!expect_ok("take batch_get_result", status, error)) { + tikv_async_request_destroy(batch_get); + tikv_async_request_destroy(scan); + tikv_async_request_destroy(reverse_scan); + tikv_async_client_destroy(client); + return 1; + } + printf("async batch_get returned %zu pairs\n", pairs.len); + for (size_t i = 0; i < pairs.len; ++i) { + printf(" get[%zu] key=", i); + print_bytes(pairs.ptr[i].key.ptr, pairs.ptr[i].key.len); + printf(" value="); + print_bytes(pairs.ptr[i].value.ptr, pairs.ptr[i].value.len); + putchar('\n'); + } + tikv_kv_pairs_destroy(pairs); + tikv_async_request_destroy(batch_get); + + pairs.ptr = NULL; + pairs.len = 0; + status = tikv_async_request_take_scan_result(scan, &pairs, &error); + if (!expect_ok("take scan_result", status, error)) { + tikv_async_request_destroy(scan); + tikv_async_request_destroy(reverse_scan); + tikv_async_client_destroy(client); + return 1; + } + printf("async scan returned %zu pairs\n", pairs.len); + for (size_t i = 0; i < pairs.len; ++i) { + printf(" scan[%zu] key=", i); + print_bytes(pairs.ptr[i].key.ptr, pairs.ptr[i].key.len); + printf(" value="); + print_bytes(pairs.ptr[i].value.ptr, pairs.ptr[i].value.len); + putchar('\n'); + } + tikv_kv_pairs_destroy(pairs); + tikv_async_request_destroy(scan); + + pairs.ptr = NULL; + pairs.len = 0; + status = tikv_async_request_take_scan_result(reverse_scan, &pairs, &error); + if (!expect_ok("take reverse_scan_result", status, error)) { + tikv_async_request_destroy(reverse_scan); + tikv_async_client_destroy(client); + return 1; + } + printf("async reverse scan returned %zu pairs\n", pairs.len); + for (size_t i = 0; i < pairs.len; ++i) { + printf(" reverse[%zu] key=", i); + print_bytes(pairs.ptr[i].key.ptr, pairs.ptr[i].key.len); + printf(" value="); + print_bytes(pairs.ptr[i].value.ptr, pairs.ptr[i].value.len); + putchar('\n'); + } + tikv_kv_pairs_destroy(pairs); + tikv_async_request_destroy(reverse_scan); + + status = tikv_async_raw_batch_write_submit( + client, cleanup_items, completion_cb, &cleanup_latch, &cleanup, &error); + if (!expect_ok("cleanup batch_write_submit", status, error) || + !wait_latch_ok("wait cleanup batch_write_submit", cleanup, &cleanup_latch)) { + tikv_async_request_destroy(cleanup); + tikv_async_client_destroy(client); + return 1; + } + tikv_async_request_destroy(cleanup); + + status = tikv_async_raw_scan_submit(client, + scan_start, + scan_end, + 16, + completion_cb, + &final_scan_latch, + &final_scan, + &error); + if (!expect_ok("final scan_submit", status, error) || + !wait_latch_ok("wait final scan_submit", final_scan, &final_scan_latch)) { + tikv_async_request_destroy(final_scan); + tikv_async_client_destroy(client); + return 1; + } + + pairs.ptr = NULL; + pairs.len = 0; + status = tikv_async_request_take_scan_result(final_scan, &pairs, &error); + if (!expect_ok("take final scan_result", status, error)) { + tikv_async_request_destroy(final_scan); + tikv_async_client_destroy(client); + return 1; + } + printf("async scan after cleanup returned %zu pairs\n", pairs.len); + tikv_kv_pairs_destroy(pairs); + tikv_async_request_destroy(final_scan); + + status = tikv_async_client_shutdown_begin(client, false, &error); + if (!expect_ok("shutdown_begin", status, error)) { + tikv_async_client_destroy(client); + return 1; + } + status = tikv_async_client_drain(client, &error); + if (!expect_ok("drain", status, error)) { + tikv_async_client_destroy(client); + return 1; + } + tikv_async_client_destroy(client); + return 0; +} diff --git a/bindings/c/examples/raw_async_demo.cpp b/bindings/c/examples/raw_async_demo.cpp new file mode 100644 index 00000000..80273665 --- /dev/null +++ b/bindings/c/examples/raw_async_demo.cpp @@ -0,0 +1,449 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#include "tikv_client_async.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using tikv::async::AsyncRequest; +using tikv::async::Client; +using tikv::async::KvPairs; +using tikv::async::Mutation; +using tikv::async::MutationWithTtl; +using tikv::async::OptionalKeyTtl; +using tikv::async::OptionalValue; +using tikv::async::Result; +using tikv::async::Slice; +using tikv::async::Status; +using tikv::async::Unit; + +bool expect_ok(const char *label, const Status &status) { + if (status.ok()) { + return true; + } + + std::cerr << label << ": " << status.to_string() << '\n'; + return false; +} + +bool parse_bool_arg(const char *value) { + return std::strcmp(value, "1") == 0 || std::strcmp(value, "true") == 0 || + std::strcmp(value, "TRUE") == 0 || + std::strcmp(value, "default-keyspace") == 0; +} + +template +bool expect_submitted(const char *label, const AsyncRequest &request) { + if (request.submitted()) { + return true; + } + + std::cerr << label << ": " << request.submit_status().to_string() << '\n'; + return false; +} + +template +bool expect_result(const char *label, const Result &result) { + if (result.ok()) { + return true; + } + + std::cerr << label << ": " << result.status().to_string() << '\n'; + return false; +} + +void print_bytes(Slice bytes) { + if (bytes.size() == 0) { + return; + } + std::cout.write(bytes.data(), static_cast(bytes.size())); +} + +void print_pairs(const char *prefix, const KvPairs &pairs) { + for (std::size_t i = 0; i < pairs.size(); ++i) { + const auto pair = pairs[i]; + std::cout << " " << prefix << '[' << i << "] key="; + print_bytes(pair.key); + std::cout << " value="; + print_bytes(pair.value); + std::cout << '\n'; + } +} + +bool wait_unit(const char *submit_label, + const char *wait_label, + AsyncRequest &request) { + return expect_submitted(submit_label, request) && + expect_ok(wait_label, request.wait()); +} + +bool read_key_ttl(Client &client, + Slice key, + const char *submit_label, + const char *wait_label, + OptionalKeyTtl &ttl) { + auto request = client.get_key_ttl_secs(key); + if (!expect_submitted(submit_label, request)) { + return false; + } + + auto result = request.wait(); + if (!expect_result(wait_label, result)) { + return false; + } + + ttl = result.take(); + return true; +} + +bool read_key_presence(Client &client, + Slice key, + const char *submit_label, + const char *wait_label, + bool &present) { + auto request = client.get(key); + if (!expect_submitted(submit_label, request)) { + return false; + } + + auto result = request.wait(); + if (!expect_result(wait_label, result)) { + return false; + } + + OptionalValue value = result.take(); + present = value.has_value(); + return true; +} + +bool wait_for_keys_absent(Client &client, + const std::vector &keys, + std::chrono::seconds timeout) { + const auto deadline = std::chrono::steady_clock::now() + timeout; + + for (;;) { + bool any_present = false; + for (Slice key : keys) { + bool present = false; + if (!read_key_presence(client, + key, + "TTL expiry get_submit", + "wait TTL expiry get_submit", + present)) { + return false; + } + any_present = any_present || present; + } + + if (!any_present) { + return true; + } + if (std::chrono::steady_clock::now() >= deadline) { + std::cerr << "TTL keys did not expire within " << timeout.count() + << " seconds\n"; + return false; + } + + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } +} + +bool expect_positive_ttl(const char *label, + const OptionalKeyTtl &ttl, + std::uint64_t configured_ttl_secs) { + if (ttl.has_value() && *ttl > 0 && *ttl <= configured_ttl_secs) { + return true; + } + + std::cerr << label << ": expected remaining TTL in [1, " + << configured_ttl_secs << "] seconds\n"; + return false; +} + +bool expect_persistent_ttl(const char *label, const OptionalKeyTtl &ttl) { + if (ttl.has_value() && *ttl == 0) { + return true; + } + + std::cerr << label << ": expected an existing key without TTL\n"; + return false; +} + +bool run_ttl_integration(Client &client) { + constexpr std::uint64_t kExpirySecs = 3; + const Slice single_expiring_key("tikv-client-rust-c/ttl/single-expiring"); + const Slice ttl_cleared_key("tikv-client-rust-c/ttl/cleared"); + const Slice batch_expiring_key("tikv-client-rust-c/ttl/batch-expiring"); + const Slice batch_persistent_key("tikv-client-rust-c/ttl/batch-persistent"); + const Slice batch_deleted_key("tikv-client-rust-c/ttl/batch-deleted"); + const Slice value("ttl-value"); + + std::cout << "running native TTL integration test\n"; + + auto put_with_ttl = + client.put_with_ttl(single_expiring_key, value, kExpirySecs); + if (!wait_unit("put_with_ttl_submit", + "wait put_with_ttl_submit", + put_with_ttl)) { + return false; + } + + OptionalKeyTtl ttl; + if (!read_key_ttl(client, + single_expiring_key, + "get_key_ttl_submit for single key", + "wait get_key_ttl_submit for single key", + ttl) || + !expect_positive_ttl("single-key TTL", ttl, kExpirySecs)) { + return false; + } + + auto put_to_clear = + client.put_with_ttl(ttl_cleared_key, value, kExpirySecs); + if (!wait_unit("put_with_ttl_submit before clearing TTL", + "wait put_with_ttl_submit before clearing TTL", + put_to_clear)) { + return false; + } + + auto clear_ttl = client.put_with_ttl(ttl_cleared_key, value, 0); + if (!wait_unit("put_with_ttl_submit clearing TTL", + "wait put_with_ttl_submit clearing TTL", + clear_ttl) || + !read_key_ttl(client, + ttl_cleared_key, + "get_key_ttl_submit after clearing TTL", + "wait get_key_ttl_submit after clearing TTL", + ttl) || + !expect_persistent_ttl("cleared TTL", ttl)) { + return false; + } + + auto seed_batch_delete = client.put(batch_deleted_key, value); + if (!wait_unit("put_submit before TTL batch delete", + "wait put_submit before TTL batch delete", + seed_batch_delete)) { + return false; + } + + auto batch_write = client.batch_write_with_ttl({ + MutationWithTtl::put(batch_expiring_key, value, kExpirySecs), + MutationWithTtl::put(batch_persistent_key, value, 0), + MutationWithTtl::remove(batch_deleted_key), + }); + if (!wait_unit("batch_write_with_ttl_submit", + "wait batch_write_with_ttl_submit", + batch_write)) { + return false; + } + + if (!read_key_ttl(client, + batch_expiring_key, + "get_key_ttl_submit for TTL batch key", + "wait get_key_ttl_submit for TTL batch key", + ttl) || + !expect_positive_ttl("batch-key TTL", ttl, kExpirySecs) || + !read_key_ttl(client, + batch_persistent_key, + "get_key_ttl_submit for persistent batch key", + "wait get_key_ttl_submit for persistent batch key", + ttl) || + !expect_persistent_ttl("persistent batch-key TTL", ttl)) { + return false; + } + + bool present = false; + if (!read_key_presence(client, + batch_deleted_key, + "get_submit for TTL batch-deleted key", + "wait get_submit for TTL batch-deleted key", + present)) { + return false; + } + if (present) { + std::cerr << "TTL batch delete did not remove its key\n"; + return false; + } + + if (!wait_for_keys_absent( + client, + {single_expiring_key, batch_expiring_key}, + std::chrono::seconds(10))) { + return false; + } + + if (!read_key_ttl(client, + single_expiring_key, + "get_key_ttl_submit after expiry", + "wait get_key_ttl_submit after expiry", + ttl)) { + return false; + } + if (ttl.has_value()) { + std::cerr << "expired key still has a TTL result\n"; + return false; + } + + for (Slice key : {ttl_cleared_key, batch_persistent_key}) { + if (!read_key_presence(client, + key, + "get_submit for persistent TTL key", + "wait get_submit for persistent TTL key", + present)) { + return false; + } + if (!present) { + std::cerr << "persistent TTL test key unexpectedly disappeared\n"; + return false; + } + } + + auto cleanup = client.batch_write({ + Mutation::remove(ttl_cleared_key), + Mutation::remove(batch_persistent_key), + }); + if (!wait_unit("cleanup TTL batch_write_submit", + "wait cleanup TTL batch_write_submit", + cleanup)) { + return false; + } + + std::cout << "native TTL integration test passed\n"; + return true; +} + +} // namespace + +int main(int argc, char **argv) { + const char *pd = argc > 1 ? argv[1] : "127.0.0.1:2379"; + const bool use_default_keyspace = argc > 2 && parse_bool_arg(argv[2]); + const bool run_ttl_tests = argc > 3 && parse_bool_arg(argv[3]); + const Slice scan_start("tikv-client-rust-c/"); + const Slice scan_end("tikv-client-rust-c0"); + const Slice key_a("tikv-client-rust-c/a"); + const Slice key_b("tikv-client-rust-c/b"); + const Slice value_a_v1("value-a-v1"); + const Slice value_a_v2("value-a-v2"); + const Slice value_b("value-b"); + + std::vector write_items = { + Mutation::put(key_a, value_a_v1), + Mutation::put(key_b, value_b), + Mutation::remove(key_a), + Mutation::put(key_a, value_a_v2), + }; + std::vector cleanup_items = { + Mutation::remove(key_a), + Mutation::remove(key_b), + }; + std::vector key_items = {key_a, key_b}; + + auto client_result = Client::connect(pd, 5, use_default_keyspace); + if (!expect_result("tikv_async_client_new", client_result)) { + return 1; + } + Client client = client_result.take(); + + auto reset = client.delete_range(scan_start, scan_end); + if (!wait_unit("delete_range_submit before batch_write", + "wait delete_range_submit before batch_write", + reset)) { + return 1; + } + + auto batch_write = client.batch_write(write_items); + if (!wait_unit("batch_write_submit", + "wait batch_write_submit", + batch_write)) { + return 1; + } + + auto batch_get = client.batch_get(key_items); + if (!expect_submitted("batch_get_submit", batch_get)) { + return 1; + } + auto scan = client.scan(scan_start, scan_end, 16); + if (!expect_submitted("scan_submit", scan)) { + return 1; + } + auto reverse_scan = client.scan_reverse(scan_start, scan_end, 16); + if (!expect_submitted("scan_reverse_submit", reverse_scan)) { + return 1; + } + + if (!expect_ok("wait batch_get_submit", batch_get.wait_status()) || + !expect_ok("wait scan_submit", scan.wait_status()) || + !expect_ok("wait scan_reverse_submit", reverse_scan.wait_status())) { + return 1; + } + + auto batch_get_result = batch_get.take(); + if (!expect_result("take batch_get_result", batch_get_result)) { + return 1; + } + KvPairs batch_get_pairs = batch_get_result.take(); + std::cout << "async batch_get returned " << batch_get_pairs.size() + << " pairs\n"; + print_pairs("get", batch_get_pairs); + + auto scan_result = scan.take(); + if (!expect_result("take scan_result", scan_result)) { + return 1; + } + KvPairs scan_pairs = scan_result.take(); + std::cout << "async scan returned " << scan_pairs.size() << " pairs\n"; + print_pairs("scan", scan_pairs); + + auto reverse_scan_result = reverse_scan.take(); + if (!expect_result("take reverse_scan_result", reverse_scan_result)) { + return 1; + } + KvPairs reverse_scan_pairs = reverse_scan_result.take(); + std::cout << "async reverse scan returned " << reverse_scan_pairs.size() + << " pairs\n"; + print_pairs("reverse", reverse_scan_pairs); + + auto cleanup = client.batch_write(cleanup_items); + if (!wait_unit("cleanup batch_write_submit", + "wait cleanup batch_write_submit", + cleanup)) { + return 1; + } + + if (run_ttl_tests && !run_ttl_integration(client)) { + return 1; + } + + auto final_scan = client.scan(scan_start, scan_end, 16); + if (!expect_submitted("final scan_submit", final_scan) || + !expect_ok("wait final scan_submit", final_scan.wait_status())) { + return 1; + } + + auto final_scan_result = final_scan.take(); + if (!expect_result("take final scan_result", final_scan_result)) { + return 1; + } + KvPairs final_pairs = final_scan_result.take(); + std::cout << "async scan after cleanup returned " << final_pairs.size() + << " pairs\n"; + + Status status = client.shutdown_begin(false); + if (!expect_ok("shutdown_begin", status)) { + return 1; + } + status = client.drain(); + if (!expect_ok("drain", status)) { + return 1; + } + + return 0; +} diff --git a/bindings/c/include/tikv_client_async.hpp b/bindings/c/include/tikv_client_async.hpp new file mode 100644 index 00000000..d301d511 --- /dev/null +++ b/bindings/c/include/tikv_client_async.hpp @@ -0,0 +1,1854 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#ifndef TIKV_CLIENT_ASYNC_HPP +#define TIKV_CLIENT_ASYNC_HPP + +#include "tikv_client_async_c.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/* + * Header-only C++17 convenience wrapper for the async C FFI in + * `tikv_client_async_c.h`. + * + * The C layer owns the actual Rust async runtime and TiKV client. This header + * adds small RAII types, `Result` error handling, and `std::future` based + * completion waiting. All submit functions still follow the C FFI contract: + * input buffers are copied by Rust before submit returns, and accepted + * requests complete exactly once. + * + * For a runnable C++ example, see + * `bindings/c/examples/raw_async_demo.cpp`. + */ +namespace tikv { +namespace async { + +/* + * Owning byte buffer used when callers want to copy TiKV bytes into C++ + * storage. + */ +using Bytes = std::vector; + +/* + * Lightweight error payload copied out of a C `tikv_error_t`. + * + * Constructing from a raw C error takes ownership of that error pointer and + * destroys it after copying the message. `Status` is usually the public type + * callers should inspect directly. + */ +struct Error { + std::string message; + + Error() = default; + + explicit Error(std::string error_message) + : message(std::move(error_message)) {} + + explicit Error(tikv_error_t *raw_error) { + if (raw_error != nullptr) { + const char *raw_message = tikv_error_message(raw_error); + if (raw_message != nullptr) { + message = raw_message; + } + tikv_error_destroy(raw_error); + } + } + + bool empty() const noexcept { return message.empty(); } +}; + +/* + * Status-only result used by operations that do not return a value. + * + * A default-constructed `Status` is OK. Non-OK statuses may carry a copied + * error message. Use `Status::from_ffi` at C boundaries so `tikv_error_t` + * ownership is handled consistently. + */ +class Status { +public: + Status() = default; + + explicit Status(tikv_status_t code) : code_(code) { + if (code_ != TIKV_STATUS_OK) { + error_.emplace(); + } + } + + Status(tikv_status_t code, std::string message) : code_(code) { + if (code_ != TIKV_STATUS_OK) { + error_.emplace(::tikv::async::Error(std::move(message))); + } + } + + Status(tikv_status_t code, ::tikv::async::Error error) : code_(code) { + if (code_ != TIKV_STATUS_OK) { + error_.emplace(std::move(error)); + } + } + + static Status OK() { return Status(); } + + static Status Error(std::string message) { + return Status(TIKV_STATUS_ERROR, std::move(message)); + } + + static Status InvalidArgument(std::string message = {}) { + return Status(TIKV_STATUS_INVALID_ARGUMENT, std::move(message)); + } + + static Status Panic(std::string message = {}) { + return Status(TIKV_STATUS_PANIC, std::move(message)); + } + + static Status Rejected(std::string message = {}) { + return Status(TIKV_STATUS_REJECTED, std::move(message)); + } + + static Status Cancelled(std::string message = {}) { + return Status(TIKV_STATUS_CANCELLED, std::move(message)); + } + + static Status from_ffi(tikv_status_t code, tikv_error_t *raw_error) { + if (code == TIKV_STATUS_OK) { + ::tikv::async::Error cleanup(raw_error); + return Status(); + } + if (raw_error == nullptr) { + return Status(code); + } + return Status(code, ::tikv::async::Error(raw_error)); + } + + bool ok() const noexcept { return code_ == TIKV_STATUS_OK; } + + tikv_status_t code() const noexcept { return code_; } + + const char *code_name() const noexcept { return status_name(code_); } + + const ::tikv::async::Error *error() const noexcept { + return error_.has_value() ? &*error_ : nullptr; + } + + const std::string &error_message() const noexcept { + static const std::string empty; + if (!error_.has_value()) { + return empty; + } + return error_->message; + } + + std::string to_string() const { + if (ok()) { + return "OK"; + } + if (error_.has_value() && !error_->message.empty()) { + std::string result(code_name()); + result += ": "; + result += error_->message; + return result; + } + return code_name(); + } + +private: + /* + * Convert a C FFI status code into a stable printable name. + */ + static const char *status_name(tikv_status_t status) noexcept { + switch (status) { + case TIKV_STATUS_OK: + return "OK"; + case TIKV_STATUS_ERROR: + return "ERROR"; + case TIKV_STATUS_INVALID_ARGUMENT: + return "INVALID_ARGUMENT"; + case TIKV_STATUS_PANIC: + return "PANIC"; + case TIKV_STATUS_REJECTED: + return "REJECTED"; + case TIKV_STATUS_CANCELLED: + return "CANCELLED"; + } + return "UNKNOWN"; + } + + tikv_status_t code_ = TIKV_STATUS_OK; + std::optional<::tikv::async::Error> error_; +}; + +/* + * Exception thrown by `Result::value()` when exceptions are enabled and the + * result is not OK. + */ +class BadResultAccess : public std::logic_error { +public: + explicit BadResultAccess(Status status) + : std::logic_error(status.to_string()), status_(std::move(status)) {} + + const Status &status() const noexcept { return status_; } + +private: + Status status_; +}; + +/* + * Marker for borrowed/view values stored in `Result`. + * + * `Result::take()` consumes an owning result value. View-like values should + * be read with `value()` so callers do not accidentally imply ownership + * transfer. + */ +struct BorrowedResultValueTag {}; + +/* + * Small `Result` type for this wrapper. + * + * Successful results contain a value. Failed results contain a non-OK + * `Status`. Check `ok()` before accessing the value, or call `value()` and + * handle `BadResultAccess` when exceptions are enabled. Use `take()` to move + * the value out after checking success. + */ +template +class Result { +public: + Result(const T &value) : value_(value) {} + + Result(T &&value) : value_(std::move(value)) {} + + Result(Status status) : status_(normalize_status(std::move(status))) {} + + bool ok() const noexcept { return status_.ok(); } + + explicit operator bool() const noexcept { return ok(); } + + const Status &status() const noexcept { return status_; } + + bool has_value() const noexcept { return value_.has_value(); } + + T &value() & { + ensure_ok(); + return *value_; + } + + const T &value() const & { + ensure_ok(); + return *value_; + } + + T &&value() && { + ensure_ok(); + return std::move(*value_); + } + + const T &&value() const && { + ensure_ok(); + return std::move(*value_); + } + + template < + typename U = T, + typename std::enable_if< + !std::is_base_of< + BorrowedResultValueTag, + typename std::remove_cv< + typename std::remove_reference::type>::type>::value, + int>::type = 0> + T take() { + ensure_ok(); + T result = std::move(*value_); + value_.reset(); + return result; + } + + T *operator->() { + return &value(); + } + + const T *operator->() const { + return &value(); + } + + T &operator*() & noexcept { + return *value_; + } + + const T &operator*() const & noexcept { + return *value_; + } + + T &&operator*() && noexcept { + return std::move(*value_); + } + +private: + static Status normalize_status(Status status) { + if (status.ok()) { + return Status::InvalidArgument("Result constructed from OK status without a value"); + } + return status; + } + + void ensure_ok() const { + if (!ok()) { + fail_bad_access(status_); + } + if (!value_.has_value()) { + fail_bad_access( + Status::InvalidArgument("Result value has already been taken")); + } + } + + [[noreturn]] static void fail_bad_access(const Status &status) { +#if defined(__cpp_exceptions) || defined(__EXCEPTIONS) || defined(_CPPUNWIND) + throw BadResultAccess(status); +#else + (void)status; + std::abort(); +#endif + } + + Status status_; + std::optional value_; +}; + +/* + * Non-owning byte slice view passed into the C FFI. + * + * `Slice` never copies bytes. It is safe to pass temporary stack buffers, + * strings, or vectors to submit functions because the Rust side copies inputs + * before submit returns. Do not keep a `Slice` after the storage it points to + * has been destroyed. + */ +class Slice { +public: + Slice() noexcept = default; + + Slice(const char *data, std::size_t size) noexcept + : data_(data), size_(size) {} + + Slice(const std::uint8_t *data, std::size_t size) noexcept + : data_(reinterpret_cast(data)), size_(size) {} + + explicit Slice(tikv_slice_t raw) noexcept + : data_(reinterpret_cast(raw.ptr)), size_(raw.len) {} + + Slice(const char *value) noexcept + : data_(value), size_(value == nullptr ? 0 : std::strlen(value)) {} + + Slice(const std::string &value) noexcept + : data_(value.data()), size_(value.size()) {} + + /* + * Deleted to avoid accidentally storing a view into a temporary string. + */ + Slice(std::string &&value) = delete; + + Slice(std::string_view value) noexcept + : data_(value.data()), size_(value.size()) {} + + tikv_slice_t to_tikv_slice() const noexcept { + return tikv_slice_t{ + reinterpret_cast(data_), + size_, + }; + } + + operator tikv_slice_t() const noexcept { return to_tikv_slice(); } + + const char *data() const noexcept { return data_; } + + const std::uint8_t *bytes() const noexcept { + return reinterpret_cast(data_); + } + + std::size_t size() const noexcept { return size_; } + + bool empty() const noexcept { return size_ == 0; } + + void clear() noexcept { + data_ = nullptr; + size_ = 0; + } + + void remove_prefix(std::size_t n) noexcept { + if (n == 0 || data_ == nullptr) { + return; + } + if (n > size_) { + n = size_; + } + data_ += n; + size_ -= n; + } + + std::string_view view_string() const noexcept { + if (data_ == nullptr || size_ == 0) { + return {}; + } + return std::string_view(data_, size_); + } + + Bytes to_bytes() const { + if (data_ == nullptr || size_ == 0) { + return {}; + } + const std::uint8_t *begin = bytes(); + return Bytes(begin, begin + size_); + } + + std::string to_string() const { + if (data_ == nullptr || size_ == 0) { + return {}; + } + return std::string(data_, size_); + } + +private: + const char *data_ = nullptr; + std::size_t size_ = 0; +}; + +/* + * RAII owner for a byte buffer returned by the Rust FFI. + * + * The buffer is released with `tikv_owned_bytes_destroy` unless callers + * explicitly transfer ownership with `release()`. + */ +class OwnedBytes { +public: + OwnedBytes() noexcept = default; + + explicit OwnedBytes(tikv_owned_bytes_t raw) noexcept : raw_(raw) {} + + ~OwnedBytes() { reset(); } + + OwnedBytes(const OwnedBytes &) = delete; + OwnedBytes &operator=(const OwnedBytes &) = delete; + + OwnedBytes(OwnedBytes &&other) noexcept : raw_(other.release()) {} + + OwnedBytes &operator=(OwnedBytes &&other) noexcept { + if (this != &other) { + reset(); + raw_ = other.release(); + } + return *this; + } + + bool empty() const noexcept { return raw_.len == 0; } + + std::size_t size() const noexcept { return raw_.len; } + + const std::uint8_t *data() const noexcept { return raw_.ptr; } + + Slice as_slice() const noexcept { return Slice(raw_.ptr, raw_.len); } + + std::string_view view_string() const noexcept { + if (raw_.ptr == nullptr || raw_.len == 0) { + return {}; + } + return std::string_view(reinterpret_cast(raw_.ptr), raw_.len); + } + + Bytes to_bytes() const { + if (raw_.ptr == nullptr || raw_.len == 0) { + return {}; + } + return Bytes(raw_.ptr, raw_.ptr + raw_.len); + } + + std::string to_string() const { + if (raw_.ptr == nullptr || raw_.len == 0) { + return {}; + } + return std::string(reinterpret_cast(raw_.ptr), raw_.len); + } + + tikv_owned_bytes_t release() noexcept { + tikv_owned_bytes_t raw = raw_; + raw_ = tikv_owned_bytes_t{nullptr, 0}; + return raw; + } + + void reset(tikv_owned_bytes_t raw = tikv_owned_bytes_t{nullptr, 0}) noexcept { + if (raw_.ptr != nullptr || raw_.len != 0) { + tikv_owned_bytes_destroy(raw_); + } + raw_ = raw; + } + + const tikv_owned_bytes_t &raw() const noexcept { return raw_; } + +private: + tikv_owned_bytes_t raw_{nullptr, 0}; +}; + +/* + * Value returned by raw get. + * + * This is move-only because it owns a Rust-allocated buffer and releases it + * with `tikv_owned_bytes_destroy`. + */ +using Value = OwnedBytes; + +/* + * Result of a raw get. + * + * `std::nullopt` means the key was not present. An engaged empty `Value` + * means the key was present and its value was the empty byte string. + */ +using OptionalValue = std::optional; + +/* + * Result of a raw get-key-TTL request. + * + * `std::nullopt` means the key was not present. An engaged value of zero means + * the key was present without a TTL. A positive value is the remaining TTL in + * seconds. + */ +using OptionalKeyTtl = std::optional; + +/* + * Marker value for operations whose successful result is only "done". + */ +struct Unit {}; + +/* + * Optional user callback invoked when an accepted request completes. + * + * The callback runs on a Rust runtime worker thread, so it must be thread-safe + * and must not throw. + */ +using CompletionCallback = void (*)(const Status &status, + void *userdata) noexcept; + +/* + * Per-request submit options. + * + * `callback` is optional; callers may instead wait on the returned + * `AsyncRequest`. `userdata` is passed through unchanged to the callback. + */ +struct SubmitOptions { + CompletionCallback callback = nullptr; + void *userdata = nullptr; +}; + +/* + * Mutation item used by `Client::batch_write`. + * + * `put` carries both key and value. `remove` carries only the key. + */ +struct Mutation { + tikv_async_raw_mutation_op_t op = TIKV_ASYNC_RAW_MUTATION_OP_PUT; + Slice key; + Slice value; + + static Mutation put(Slice key, Slice value) noexcept { + return Mutation{TIKV_ASYNC_RAW_MUTATION_OP_PUT, key, value}; + } + + static Mutation remove(Slice key) noexcept { + return Mutation{TIKV_ASYNC_RAW_MUTATION_OP_DELETE, key, Slice()}; + } +}; + +/* + * TTL-capable mutation item used by `Client::batch_write_with_ttl`. + * + * PUT requires an explicit relative TTL in seconds; zero disables TTL for the + * key. DELETE ignores TTL and always carries zero. + */ +struct MutationWithTtl { + tikv_async_raw_mutation_op_t op = TIKV_ASYNC_RAW_MUTATION_OP_PUT; + Slice key; + Slice value; + std::uint64_t ttl_secs = 0; + + static MutationWithTtl put(Slice key, + Slice value, + std::uint64_t ttl_secs) noexcept { + return MutationWithTtl{ + TIKV_ASYNC_RAW_MUTATION_OP_PUT, + key, + value, + ttl_secs, + }; + } + + static MutationWithTtl remove(Slice key) noexcept { + return MutationWithTtl{ + TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + key, + Slice(), + 0, + }; + } +}; + +/* + * Borrowed key/value pair view. + * + * The slices point into the owning `KvPairs` object and become invalid when + * that object is moved, reset, or destroyed. + */ +struct KvPairView { + Slice key; + Slice value; +}; + +/* + * RAII owner for a vector of key/value pairs returned by scan or batch-get. + * + * Iterators and `KvPairView` values borrow from this object. The underlying + * Rust-owned memory is released with `tikv_kv_pairs_destroy`. + */ +class KvPairs { +public: + class const_iterator { + public: + using iterator_category = std::forward_iterator_tag; + using value_type = KvPairView; + using difference_type = std::ptrdiff_t; + using pointer = void; + using reference = KvPairView; + + const_iterator() noexcept = default; + + explicit const_iterator(const tikv_kv_pair_t *ptr) noexcept : ptr_(ptr) {} + + KvPairView operator*() const noexcept { return KvPairs::pair_view(*ptr_); } + + const_iterator &operator++() noexcept { + ++ptr_; + return *this; + } + + const_iterator operator++(int) noexcept { + const_iterator current(*this); + ++(*this); + return current; + } + + bool operator==(const const_iterator &other) const noexcept { + return ptr_ == other.ptr_; + } + + bool operator!=(const const_iterator &other) const noexcept { + return !(*this == other); + } + + private: + const tikv_kv_pair_t *ptr_ = nullptr; + }; + + KvPairs() noexcept = default; + + explicit KvPairs(tikv_kv_pairs_t raw) noexcept : raw_(raw) {} + + ~KvPairs() { reset(); } + + KvPairs(const KvPairs &) = delete; + KvPairs &operator=(const KvPairs &) = delete; + + KvPairs(KvPairs &&other) noexcept : raw_(other.release()) {} + + KvPairs &operator=(KvPairs &&other) noexcept { + if (this != &other) { + reset(); + raw_ = other.release(); + } + return *this; + } + + bool empty() const noexcept { return raw_.len == 0; } + + std::size_t size() const noexcept { return raw_.len; } + + KvPairView operator[](std::size_t index) const noexcept { + return pair_view(raw_.ptr[index]); + } + + const_iterator begin() const noexcept { return const_iterator(raw_.ptr); } + + const_iterator end() const noexcept { + return const_iterator(raw_.ptr == nullptr ? nullptr : raw_.ptr + raw_.len); + } + + tikv_kv_pairs_t release() noexcept { + tikv_kv_pairs_t raw = raw_; + raw_ = tikv_kv_pairs_t{nullptr, 0}; + return raw; + } + + void reset(tikv_kv_pairs_t raw = tikv_kv_pairs_t{nullptr, 0}) noexcept { + if (raw_.ptr != nullptr || raw_.len != 0) { + tikv_kv_pairs_destroy(raw_); + } + raw_ = raw; + } + + const tikv_kv_pairs_t &raw() const noexcept { return raw_; } + +private: + static Slice bytes_view(const tikv_owned_bytes_t &bytes) noexcept { + return Slice(bytes.ptr, bytes.len); + } + + static KvPairView pair_view(const tikv_kv_pair_t &pair) noexcept { + return KvPairView{ + bytes_view(pair.key), + bytes_view(pair.value), + }; + } + + tikv_kv_pairs_t raw_{nullptr, 0}; +}; + +/* + * RAII owner for an accepted async request handle. + * + * A `Request` remains valid until destroyed or reset. Result extraction + * helpers must be called only after successful completion, and each successful + * result may be taken only once. + */ +class Request { +public: + Request() noexcept = default; + + explicit Request(tikv_async_request_t *raw) noexcept : raw_(raw) {} + + ~Request() { reset(); } + + Request(const Request &) = delete; + Request &operator=(const Request &) = delete; + + Request(Request &&other) noexcept : raw_(other.release()) {} + + Request &operator=(Request &&other) noexcept { + if (this != &other) { + reset(); + raw_ = other.release(); + } + return *this; + } + + bool valid() const noexcept { return raw_ != nullptr; } + + explicit operator bool() const noexcept { return valid(); } + + /* + * Return the borrowed raw C handle. The `Request` keeps ownership. + */ + tikv_async_request_t *get() const noexcept { return raw_; } + + /* + * Transfer ownership of the raw C handle to the caller. + */ + tikv_async_request_t *release() noexcept { + tikv_async_request_t *raw = raw_; + raw_ = nullptr; + return raw; + } + + /* + * Replace the owned request handle, destroying the previous one if present. + */ + void reset(tikv_async_request_t *raw = nullptr) noexcept { + if (raw_ != nullptr) { + tikv_async_request_destroy(raw_); + } + raw_ = raw; + } + + /* + * Request best-effort cancellation. + * + * A successful return value of `true` means this call newly requested + * cancellation. `false` means the request had already completed or had + * already been asked to cancel. + */ + Result cancel() { + if (raw_ == nullptr) { + return invalid_request_status(); + } + + bool cancel_requested = false; + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_async_request_cancel(raw_, &cancel_requested, &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + return status; + } + return cancel_requested; + } + + /* + * Query whether the request callback has completed. + */ + Result is_complete() const { + if (raw_ == nullptr) { + return invalid_request_status(); + } + + bool complete = false; + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_async_request_is_complete(raw_, &complete, &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + return status; + } + return complete; + } + + /* + * Return the terminal completion status. + * + * This is valid only after the request has completed. + */ + Result completion_status() const { + if (raw_ == nullptr) { + return invalid_request_status(); + } + + tikv_status_t completion = TIKV_STATUS_ERROR; + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_async_request_completion_status(raw_, &completion, &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + return status; + } + return completion; + } + + /* + * Copy the terminal error message out of the request state, if any. + */ + std::string error_message() const { + std::string_view message = error_message_view(); + if (message.empty()) { + return {}; + } + return std::string(message.data(), message.size()); + } + + /* + * View the terminal error message without copying it. + * + * The view remains valid only while this `Request` owns the raw handle. + */ + std::string_view error_message_view() const noexcept { + const char *message = error_c_str(); + if (message == nullptr) { + return {}; + } + return std::string_view(message, std::strlen(message)); + } + + /* + * Take the result from a completed raw get request. + * + * `std::nullopt` means the key was not present. Calling this twice on the + * same successful request returns an error. + */ + Result take_get_result() { + if (raw_ == nullptr) { + return invalid_request_status(); + } + + tikv_owned_bytes_t raw_value{nullptr, 0}; + bool found = false; + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_async_request_take_get_result(raw_, &raw_value, &found, &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + tikv_owned_bytes_destroy(raw_value); + return status; + } + + if (!found) { + tikv_owned_bytes_destroy(raw_value); + return Result(OptionalValue(std::nullopt)); + } + + return Result(OptionalValue(std::in_place, raw_value)); + } + + /* + * Take the result from a completed raw get-key-TTL request. + * + * `std::nullopt` means the key was not present. An engaged zero means the key + * exists without a TTL. Calling this twice on the same successful request + * returns an error. + */ + Result take_get_key_ttl_result() { + if (raw_ == nullptr) { + return invalid_request_status(); + } + + std::uint64_t ttl_secs = 0; + bool found = false; + tikv_error_t *error = nullptr; + const tikv_status_t code = tikv_async_request_take_get_key_ttl_result( + raw_, &ttl_secs, &found, &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + return status; + } + + if (!found) { + return Result(OptionalKeyTtl(std::nullopt)); + } + + return Result(OptionalKeyTtl(ttl_secs)); + } + + /* + * Take the result from a completed scan or reverse-scan request. + */ + Result take_scan_result() { + return take_pairs_result(tikv_async_request_take_scan_result); + } + + /* + * Take the result from a completed batch-get request. + */ + Result take_batch_get_result() { + return take_pairs_result(tikv_async_request_take_batch_get_result); + } + +private: + using TakePairsFn = std::function; + + const char *error_c_str() const noexcept { + if (raw_ == nullptr) { + return nullptr; + } + return tikv_async_request_error_message(raw_); + } + + static Status invalid_request_status() { + return Status::InvalidArgument("request is null"); + } + + Result take_pairs_result(const TakePairsFn &take_pairs) { + if (raw_ == nullptr) { + return invalid_request_status(); + } + if (!take_pairs) { + return Status::InvalidArgument("take pairs function is empty"); + } + + tikv_kv_pairs_t raw_pairs{nullptr, 0}; + tikv_error_t *error = nullptr; + const tikv_status_t code = take_pairs(raw_, &raw_pairs, &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + tikv_kv_pairs_destroy(raw_pairs); + return status; + } + return Result(KvPairs(raw_pairs)); + } + + tikv_async_request_t *raw_ = nullptr; +}; + +/* + * Submitted async request that eventually yields a value of type `T`. + * + * `submitted()` tells whether the submit call itself was accepted. Once a + * request is accepted, `wait_status()` waits for terminal completion and + * `take()` extracts the result. For convenience, `wait()` does both. + */ +template +class AsyncRequest { +public: + /* + * Converts a completed raw request handle into the typed result. + */ + using Extractor = std::function(Request &)>; + + AsyncRequest() = default; + + /* + * Build a request object for a submit call that was rejected synchronously. + */ + static AsyncRequest rejected(Status status) { + AsyncRequest request; + request.submit_status_ = normalize_rejected_status(std::move(status)); + return request; + } + + /* + * Build a request object for a submit call accepted by the C FFI. + */ + static AsyncRequest submitted(Request request, + std::shared_future done, + Extractor extractor) { + AsyncRequest async_request; + async_request.submit_status_ = Status::OK(); + async_request.request_ = std::move(request); + async_request.done_ = std::move(done); + async_request.extractor_ = std::move(extractor); + return async_request; + } + + AsyncRequest(const AsyncRequest &) = delete; + AsyncRequest &operator=(const AsyncRequest &) = delete; + + AsyncRequest(AsyncRequest &&) noexcept = default; + AsyncRequest &operator=(AsyncRequest &&) noexcept = default; + + bool submitted() const noexcept { return submit_status_.ok(); } + + /* + * Status returned by the submit call itself. + * + * A non-OK status means no async callback will run. + */ + const Status &submit_status() const noexcept { return submit_status_; } + + /* + * Non-blocking completion check. + */ + bool ready() const { + if (!submitted()) { + return true; + } + if (!done_.valid()) { + return true; + } + return done_.wait_for(std::chrono::seconds(0)) == std::future_status::ready; + } + + /* + * Wait for completion and return only the terminal status. + */ + Status wait_status() const { + if (!submit_status_.ok()) { + return submit_status_; + } + if (!done_.valid()) { + return Status::InvalidArgument("async request has no completion signal"); + } + return done_.get(); + } + + /* + * Extract the typed result after successful completion. + * + * The underlying C result can be taken only once. + */ + Result take() { + if (!submit_status_.ok()) { + return submit_status_; + } + if (!extractor_) { + return Status::InvalidArgument("async request has no result extractor"); + } + return extractor_(request_); + } + + /* + * Wait for completion, then take the typed result. + */ + Result wait() { + Status status = wait_status(); + if (!status.ok()) { + return status; + } + return take(); + } + + /* + * Request best-effort cancellation of the accepted request. + */ + Result cancel() { + if (!submit_status_.ok()) { + return submit_status_; + } + return request_.cancel(); + } + + /* + * Access the underlying request handle for lower-level operations. + */ + Request &request() noexcept { return request_; } + + const Request &request() const noexcept { return request_; } + +private: + static Status normalize_rejected_status(Status status) { + if (status.ok()) { + return Status::InvalidArgument("async request rejected with OK status"); + } + return status; + } + + Status submit_status_ = Status::InvalidArgument("async request is not submitted"); + Request request_; + std::shared_future done_; + Extractor extractor_; +}; + +/* + * `AsyncRequest` specialization for operations that return no payload. + * + * `wait()` returns the terminal `Status` directly because there is no value to + * extract after completion. + */ +template <> +class AsyncRequest { +public: + AsyncRequest() = default; + + static AsyncRequest rejected(Status status) { + AsyncRequest request; + request.submit_status_ = normalize_rejected_status(std::move(status)); + return request; + } + + static AsyncRequest submitted(Request request, std::shared_future done) { + AsyncRequest async_request; + async_request.submit_status_ = Status::OK(); + async_request.request_ = std::move(request); + async_request.done_ = std::move(done); + return async_request; + } + + AsyncRequest(const AsyncRequest &) = delete; + AsyncRequest &operator=(const AsyncRequest &) = delete; + + AsyncRequest(AsyncRequest &&) noexcept = default; + AsyncRequest &operator=(AsyncRequest &&) noexcept = default; + + bool submitted() const noexcept { return submit_status_.ok(); } + + /* + * Status returned by the submit call itself. + */ + const Status &submit_status() const noexcept { return submit_status_; } + + /* + * Non-blocking completion check. + */ + bool ready() const { + if (!submitted()) { + return true; + } + if (!done_.valid()) { + return true; + } + return done_.wait_for(std::chrono::seconds(0)) == std::future_status::ready; + } + + /* + * Wait for terminal completion. + */ + Status wait() const { + if (!submit_status_.ok()) { + return submit_status_; + } + if (!done_.valid()) { + return Status::InvalidArgument("async request has no completion signal"); + } + return done_.get(); + } + + /* + * Request best-effort cancellation of the accepted request. + */ + Result cancel() { + if (!submit_status_.ok()) { + return submit_status_; + } + return request_.cancel(); + } + + /* + * Access the underlying request handle for lower-level operations. + */ + Request &request() noexcept { return request_; } + + const Request &request() const noexcept { return request_; } + +private: + static Status normalize_rejected_status(Status status) { + if (status.ok()) { + return Status::InvalidArgument("async request rejected with OK status"); + } + return status; + } + + Status submit_status_ = Status::InvalidArgument("async request is not submitted"); + Request request_; + std::shared_future done_; +}; + +/* + * RAII owner for an async TiKV raw client. + * + * Destroying a `Client` begins shutdown with cancellation, drains accepted + * requests, and releases the C client handle. Call `shutdown_begin(false)` and + * `drain()` explicitly if you want in-flight requests to finish normally. + */ +class Client { +public: + /* + * Shared C++ owner type for callers that want many application objects to + * hold the same client instance. The C ABI handle is still owned by exactly + * one `Client` object and is destroyed when the last shared owner releases it. + */ + using SharedPtr = std::shared_ptr; + + Client() noexcept = default; + + explicit Client(tikv_async_client_t *raw) noexcept : raw_(raw) {} + + /* + * Connect to one or more PD endpoints. + * + * `pd_endpoints_csv` is a comma-separated endpoint list such as + * `"127.0.0.1:2379,127.0.0.1:2381"`. `timeout_secs == 0` uses the Rust + * client's default timeout. Set `use_default_keyspace` for API v2 clusters + * that require the default keyspace. By default, the gRPC response decoding + * size is unlimited for compatibility. Pass a byte limit explicitly to + * override it; zero uses the Rust client's 4 MiB default. + */ + static Result connect(const char *pd_endpoints_csv, + std::uint32_t timeout_secs = 0, + bool use_default_keyspace = false, + std::size_t grpc_max_decoding_message_size = + std::numeric_limits::max()) { + tikv_async_client_t *raw_client = nullptr; + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_async_client_new_with_grpc_max_decoding_message_size( + pd_endpoints_csv, + timeout_secs, + use_default_keyspace, + grpc_max_decoding_message_size, + &raw_client, + &error); + Status status = Status::from_ffi(code, error); + if (!status.ok()) { + if (raw_client != nullptr) { + tikv_async_client_destroy(raw_client); + } + return status; + } + return Client(raw_client); + } + + /* + * String overload for `connect`. + */ + static Result connect(const std::string &pd_endpoints_csv, + std::uint32_t timeout_secs = 0, + bool use_default_keyspace = false, + std::size_t grpc_max_decoding_message_size = + std::numeric_limits::max()) { + return connect(pd_endpoints_csv.c_str(), + timeout_secs, + use_default_keyspace, + grpc_max_decoding_message_size); + } + + /* + * Try to connect and return the resulting client inside a `std::shared_ptr`. + * + * This is a convenience factory for dependency injection or worker objects + * that should share one C++ `Client`. It does not clone the C ABI handle. + */ + static Result connect_as_shared(const char *pd_endpoints_csv, + std::uint32_t timeout_secs = 0, + bool use_default_keyspace = false, + std::size_t grpc_max_decoding_message_size = + std::numeric_limits::max()) { + auto client_result = connect(pd_endpoints_csv, + timeout_secs, + use_default_keyspace, + grpc_max_decoding_message_size); + if (!client_result.ok()) { + return client_result.status(); + } + + return std::make_shared(client_result.take()); + } + + /* + * String overload for `connect_as_shared`. + */ + static Result connect_as_shared(const std::string &pd_endpoints_csv, + std::uint32_t timeout_secs = 0, + bool use_default_keyspace = false, + std::size_t grpc_max_decoding_message_size = + std::numeric_limits::max()) { + return connect_as_shared(pd_endpoints_csv.c_str(), + timeout_secs, + use_default_keyspace, + grpc_max_decoding_message_size); + } + + ~Client() { reset(); } + + Client(const Client &) = delete; + Client &operator=(const Client &) = delete; + + Client(Client &&other) noexcept : raw_(other.release()) {} + + Client &operator=(Client &&other) noexcept { + if (this != &other) { + reset(); + raw_ = other.release(); + } + return *this; + } + + bool valid() const noexcept { return raw_ != nullptr; } + + explicit operator bool() const noexcept { return valid(); } + + /* + * Return the borrowed raw C client handle. The `Client` keeps ownership. + */ + tikv_async_client_t *get() const noexcept { return raw_; } + + /* + * Transfer ownership of the raw C client handle to the caller. + */ + tikv_async_client_t *release() noexcept { + tikv_async_client_t *raw = raw_; + raw_ = nullptr; + return raw; + } + + /* + * Replace the owned client handle, destroying the previous one if present. + */ + void reset(tikv_async_client_t *raw = nullptr) noexcept { + if (raw_ != nullptr) { + tikv_async_client_destroy(raw_); + } + raw_ = raw; + } + + /* + * Start client shutdown. + * + * After shutdown begins, new submit calls are rejected. If + * `cancel_inflight` is true, accepted in-flight requests complete as + * cancelled. + */ + Status shutdown_begin(bool cancel_inflight = false) { + if (raw_ == nullptr) { + return invalid_client_status(); + } + + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_async_client_shutdown_begin(raw_, cancel_inflight, &error); + return Status::from_ffi(code, error); + } + + /* + * Wait for all currently accepted requests to finish their callbacks. + */ + Status drain() const { + if (raw_ == nullptr) { + return invalid_client_status(); + } + + tikv_error_t *error = nullptr; + const tikv_status_t code = tikv_async_client_drain(raw_, &error); + return Status::from_ffi(code, error); + } + + /* + * Submit a raw put request. + * + * The key and value bytes are copied by Rust before this function returns. + */ + AsyncRequest put(Slice key, + Slice value, + SubmitOptions options = SubmitOptions()) const { + return submit_unit( + [this, key, value](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_put_submit(raw_, + key.to_tikv_slice(), + value.to_tikv_slice(), + completion, + userdata, + out_request, + out_error); + }, + options); + } + + /* + * Submit a raw put with a relative TTL in seconds. + * + * `ttl_secs == 0` disables TTL for the key. The key and value bytes are + * copied by Rust before this function returns. + */ + AsyncRequest put_with_ttl( + Slice key, + Slice value, + std::uint64_t ttl_secs, + SubmitOptions options = SubmitOptions()) const { + return submit_unit( + [this, key, value, ttl_secs](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_put_with_ttl_submit(raw_, + key.to_tikv_slice(), + value.to_tikv_slice(), + ttl_secs, + completion, + userdata, + out_request, + out_error); + }, + options); + } + + /* + * Submit a raw get request. + * + * Use `wait()` or `wait_status()` followed by `take()` to retrieve an + * `OptionalValue`. + */ + AsyncRequest get(Slice key, + SubmitOptions options = SubmitOptions()) const { + return submit_value( + [this, key](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_get_submit(raw_, + key.to_tikv_slice(), + completion, + userdata, + out_request, + out_error); + }, + [](Request &request) { return request.take_get_result(); }, + options); + } + + /* + * Submit a request for the remaining raw key TTL in seconds. + * + * This method returns an `AsyncRequest`. Call `wait()` on that request to + * obtain a `Result`. + */ + AsyncRequest get_key_ttl_secs( + Slice key, + SubmitOptions options = SubmitOptions()) const { + return submit_value( + [this, key](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_get_key_ttl_submit(raw_, + key.to_tikv_slice(), + completion, + userdata, + out_request, + out_error); + }, + [](Request &request) { return request.take_get_key_ttl_result(); }, + options); + } + + /* + * Submit a raw delete request for one key. + */ + AsyncRequest remove(Slice key, + SubmitOptions options = SubmitOptions()) const { + return submit_unit( + [this, key](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_delete_submit(raw_, + key.to_tikv_slice(), + completion, + userdata, + out_request, + out_error); + }, + options); + } + + /* + * Submit a raw delete-range request. + * + * An empty `end_key` means the range is unbounded above. + */ + AsyncRequest delete_range(Slice start_key, + Slice end_key, + SubmitOptions options = SubmitOptions()) const { + return submit_unit( + [this, start_key, end_key](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_delete_range_submit(raw_, + start_key.to_tikv_slice(), + end_key.to_tikv_slice(), + completion, + userdata, + out_request, + out_error); + }, + options); + } + + /* + * Submit a forward raw scan over `[start_key, end_key)`. + * + * `limit` caps the number of returned pairs. An empty `end_key` means the + * range is unbounded above. + */ + AsyncRequest scan(Slice start_key, + Slice end_key, + std::uint32_t limit, + SubmitOptions options = SubmitOptions()) const { + return submit_value( + [this, start_key, end_key, limit](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_scan_submit(raw_, + start_key.to_tikv_slice(), + end_key.to_tikv_slice(), + limit, + completion, + userdata, + out_request, + out_error); + }, + [](Request &request) { return request.take_scan_result(); }, + options); + } + + /* + * Submit a reverse raw scan over `[start_key, end_key)`. + * + * Returned pairs are in reverse lexicographical key order. The underlying + * Rust raw client currently does not support reverse scanning from an empty + * start key. + */ + AsyncRequest scan_reverse(Slice start_key, + Slice end_key, + std::uint32_t limit, + SubmitOptions options = SubmitOptions()) const { + return submit_value( + [this, start_key, end_key, limit](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_scan_reverse_submit(raw_, + start_key.to_tikv_slice(), + end_key.to_tikv_slice(), + limit, + completion, + userdata, + out_request, + out_error); + }, + [](Request &request) { return request.take_scan_result(); }, + options); + } + + /* + * Submit a raw batch-get request. + * + * Only existing keys are returned, and TiKV does not guarantee that result + * order matches input key order. + */ + AsyncRequest batch_get(const std::vector &keys, + SubmitOptions options = SubmitOptions()) const { + std::vector raw_keys; + raw_keys.reserve(keys.size()); + for (const Slice &key : keys) { + raw_keys.push_back(key.to_tikv_slice()); + } + + const tikv_slices_t ffi_keys{raw_keys.empty() ? nullptr : raw_keys.data(), + raw_keys.size()}; + /* + * The C submit call synchronously copies `ffi_keys`, so it is safe for this + * vector-backed view to die after `submit_value` returns. + */ + return submit_value( + [this, ffi_keys](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_batch_get_submit(raw_, + ffi_keys, + completion, + userdata, + out_request, + out_error); + }, + [](Request &request) { return request.take_batch_get_result(); }, + options); + } + + /* + * Initializer-list overload for `batch_get`. + */ + AsyncRequest batch_get(std::initializer_list keys, + SubmitOptions options = SubmitOptions()) const { + return batch_get(std::vector(keys), options); + } + + /* + * Submit a mixed raw write batch. + * + * The C/Rust layer preserves put/delete phase order while coalescing + * adjacent compatible mutations into efficient lower-level calls. + */ + AsyncRequest batch_write(const std::vector &mutations, + SubmitOptions options = SubmitOptions()) const { + std::vector raw_mutations; + raw_mutations.reserve(mutations.size()); + for (const Mutation &mutation : mutations) { + raw_mutations.push_back(tikv_async_raw_mutation_t{ + mutation.op, + mutation.key.to_tikv_slice(), + mutation.value.to_tikv_slice(), + }); + } + + const tikv_async_raw_mutations_t ffi_mutations{ + raw_mutations.empty() ? nullptr : raw_mutations.data(), + raw_mutations.size(), + }; + /* + * The C submit call synchronously copies `ffi_mutations`, so it is safe for + * this vector-backed view to die after `submit_unit` returns. + */ + return submit_unit( + [this, ffi_mutations](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_batch_write_submit(raw_, + ffi_mutations, + completion, + userdata, + out_request, + out_error); + }, + options); + } + + /* + * Initializer-list overload for `batch_write`. + */ + AsyncRequest batch_write(std::initializer_list mutations, + SubmitOptions options = SubmitOptions()) const { + return batch_write(std::vector(mutations), options); + } + + /* + * Submit a TTL-capable mixed raw write batch. + * + * Each PUT carries its own relative TTL in seconds. DELETE ignores TTL. The + * C submit call synchronously copies all mutation inputs before this method + * returns. + */ + AsyncRequest batch_write_with_ttl( + const std::vector &mutations, + SubmitOptions options = SubmitOptions()) const { + std::vector raw_mutations; + raw_mutations.reserve(mutations.size()); + for (const MutationWithTtl &mutation : mutations) { + raw_mutations.push_back(tikv_async_raw_mutation_with_ttl_t{ + mutation.op, + mutation.key.to_tikv_slice(), + mutation.value.to_tikv_slice(), + mutation.ttl_secs, + }); + } + + const tikv_async_raw_mutations_with_ttl_t ffi_mutations{ + raw_mutations.empty() ? nullptr : raw_mutations.data(), + raw_mutations.size(), + }; + return submit_unit( + [this, ffi_mutations](tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error) { + return tikv_async_raw_batch_write_with_ttl_submit(raw_, + ffi_mutations, + completion, + userdata, + out_request, + out_error); + }, + options); + } + + /* + * Initializer-list overload for `batch_write_with_ttl`. + */ + AsyncRequest batch_write_with_ttl( + std::initializer_list mutations, + SubmitOptions options = SubmitOptions()) const { + return batch_write_with_ttl(std::vector(mutations), + options); + } + +private: + /* + * Heap state handed to the C callback. + * + * The callback takes ownership with `std::unique_ptr`, fulfills the future, + * then invokes the optional user callback. + */ + struct RequestCompletionState { + std::promise completion; + SubmitOptions options; + }; + + using SubmitFn = std::function; + + static Status invalid_client_status() { + return Status::InvalidArgument("client is null"); + } + + static Status build_completion_status(tikv_async_request_t *request, + tikv_status_t final_status) { + if (final_status == TIKV_STATUS_OK) { + return Status::OK(); + } + + const char *message = tikv_async_request_error_message(request); + if (message != nullptr) { + return Status(final_status, std::string(message)); + } + return Status(final_status); + } + + /* + * C callback bridge used for every request submitted through this wrapper. + */ + static void handle_request_completion(tikv_async_request_t *request, + tikv_status_t final_status, + void *userdata) noexcept { + std::unique_ptr state( + static_cast(userdata)); + + Status completion_status = build_completion_status(request, final_status); + state->completion.set_value(completion_status); + + if (state->options.callback != nullptr) { + state->options.callback(completion_status, state->options.userdata); + } + } + + /* + * Common submit path for requests whose successful completion yields a + * payload. + */ + template + AsyncRequest submit_value(SubmitFn submit, + typename AsyncRequest::Extractor extractor, + SubmitOptions options) const { + if (raw_ == nullptr) { + return AsyncRequest::rejected(invalid_client_status()); + } + if (!submit) { + return AsyncRequest::rejected(Status::InvalidArgument("submit function is empty")); + } + + RequestCompletionState *state = new RequestCompletionState(); + state->options = options; + std::shared_future done = state->completion.get_future().share(); + + tikv_async_request_t *raw_request = nullptr; + tikv_error_t *error = nullptr; + const tikv_status_t code = + submit(handle_request_completion, state, &raw_request, &error); + Status submit_status = Status::from_ffi(code, error); + if (!submit_status.ok()) { + delete state; + return AsyncRequest::rejected(std::move(submit_status)); + } + + return AsyncRequest::submitted(Request(raw_request), + std::move(done), + std::move(extractor)); + } + + /* + * Common submit path for requests whose successful completion has no + * payload. + */ + AsyncRequest submit_unit(SubmitFn submit, SubmitOptions options) const { + if (raw_ == nullptr) { + return AsyncRequest::rejected(invalid_client_status()); + } + if (!submit) { + return AsyncRequest::rejected(Status::InvalidArgument("submit function is empty")); + } + + RequestCompletionState *state = new RequestCompletionState(); + state->options = options; + std::shared_future done = state->completion.get_future().share(); + + tikv_async_request_t *raw_request = nullptr; + tikv_error_t *error = nullptr; + const tikv_status_t code = + submit(handle_request_completion, state, &raw_request, &error); + Status submit_status = Status::from_ffi(code, error); + if (!submit_status.ok()) { + delete state; + return AsyncRequest::rejected(std::move(submit_status)); + } + + return AsyncRequest::submitted(Request(raw_request), std::move(done)); + } + + tikv_async_client_t *raw_ = nullptr; +}; + +} // namespace async +} // namespace tikv + +#endif // TIKV_CLIENT_ASYNC_HPP diff --git a/bindings/c/include/tikv_client_async_c.h b/bindings/c/include/tikv_client_async_c.h new file mode 100644 index 00000000..1e45621b --- /dev/null +++ b/bindings/c/include/tikv_client_async_c.h @@ -0,0 +1,502 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#ifndef TIKV_CLIENT_ASYNC_C_H +#define TIKV_CLIENT_ASYNC_C_H + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * Async C ABI contract. + * + * Design goals for this binding: + * + * - Expose a natural async ABI: submit + completion callback. + * - Copy every cross-async-boundary input before submit returns. + * - Make shutdown/drain/cancel semantics explicit and defensive. + * + * Threading model: + * + * - Submit functions may be called concurrently on the same client handle. + * - Completion callbacks run on Rust async runtime worker threads, not on the + * submitter thread. + * - Completion callbacks must therefore be thread-safe. + * + * Completion contract: + * + * - Every request accepted by a submit function triggers exactly one completion + * callback, including cancellation and shutdown-triggered cancellation. + * - A submit function returning non-`TIKV_STATUS_OK` means the request was not + * accepted and no completion callback will run. + * + * Request handle contract: + * + * - Submit success returns a request handle. + * - The request handle remains valid until `tikv_async_request_destroy`. + * - It is valid to destroy the request handle before the callback fires; Rust + * keeps the internal state alive until the callback returns. + * - Result accessors may be called only after completion. + * - Returned buffers and errors are Rust-owned and must be released with the + * matching destroy helper from this header. + * + * Client shutdown contract: + * + * - `tikv_async_client_shutdown_begin` immediately rejects new submit calls. + * - `tikv_async_client_shutdown_begin(..., false, ...)` lets in-flight requests + * finish normally. + * - `tikv_async_client_shutdown_begin(..., true, ...)` requests cancellation of + * every in-flight request; each accepted request still completes through its + * callback with `TIKV_STATUS_CANCELLED`. + * - `tikv_async_client_drain` blocks until every accepted in-flight request has + * finished its callback. + * - `tikv_async_client_destroy` performs shutdown-begin with cancellation, + * drains, and then releases the client. + * - Do not call `tikv_async_client_drain` or `tikv_async_client_destroy` from a + * completion callback belonging to the same client, or you may deadlock. + */ + +typedef struct tikv_async_client_t tikv_async_client_t; +typedef struct tikv_async_request_t tikv_async_request_t; +typedef struct tikv_error_t tikv_error_t; + +typedef struct tikv_slice_t { + const uint8_t *ptr; + size_t len; +} tikv_slice_t; + +typedef struct tikv_slices_t { + const tikv_slice_t *ptr; + size_t len; +} tikv_slices_t; + +typedef enum tikv_async_raw_mutation_op_t { + TIKV_ASYNC_RAW_MUTATION_OP_PUT = 0, + TIKV_ASYNC_RAW_MUTATION_OP_DELETE = 1 +} tikv_async_raw_mutation_op_t; + +typedef struct tikv_async_raw_mutation_t { + tikv_async_raw_mutation_op_t op; + tikv_slice_t key; + /* + * Used only for PUT. + * DELETE ignores this field, so callers may pass {NULL, 0}. + */ + tikv_slice_t value; +} tikv_async_raw_mutation_t; + +typedef struct tikv_async_raw_mutations_t { + const tikv_async_raw_mutation_t *ptr; + size_t len; +} tikv_async_raw_mutations_t; + +/* + * TTL-capable mutation used by + * `tikv_async_raw_batch_write_with_ttl_submit`. + * + * `ttl_secs` is a relative TTL in seconds and is used only for PUT. A value of + * 0 disables TTL for the key. DELETE ignores both `value` and `ttl_secs`. + */ +typedef struct tikv_async_raw_mutation_with_ttl_t { + tikv_async_raw_mutation_op_t op; + tikv_slice_t key; + tikv_slice_t value; + uint64_t ttl_secs; +} tikv_async_raw_mutation_with_ttl_t; + +typedef struct tikv_async_raw_mutations_with_ttl_t { + const tikv_async_raw_mutation_with_ttl_t *ptr; + size_t len; +} tikv_async_raw_mutations_with_ttl_t; + +typedef struct tikv_owned_bytes_t { + uint8_t *ptr; + size_t len; +} tikv_owned_bytes_t; + +typedef struct tikv_kv_pair_t { + tikv_owned_bytes_t key; + tikv_owned_bytes_t value; +} tikv_kv_pair_t; + +typedef struct tikv_kv_pairs_t { + tikv_kv_pair_t *ptr; + size_t len; +} tikv_kv_pairs_t; + +typedef enum tikv_status_t { + TIKV_STATUS_OK = 0, + TIKV_STATUS_ERROR = 1, + TIKV_STATUS_INVALID_ARGUMENT = 2, + TIKV_STATUS_PANIC = 3, + TIKV_STATUS_REJECTED = 4, + TIKV_STATUS_CANCELLED = 5 +} tikv_status_t; + +/* + * Completion callback type. + * + * `final_status` is the terminal status of the accepted request: + * + * - `TIKV_STATUS_OK`: request succeeded. + * - `TIKV_STATUS_ERROR`: request failed; inspect + * `tikv_async_request_error_message`. + * - `TIKV_STATUS_CANCELLED`: request was cancelled explicitly or by client + * shutdown; inspect `tikv_async_request_error_message` for details. + * - `TIKV_STATUS_PANIC`: Rust caught a panic while executing the request. + * + * The callback may call request result accessors, `tikv_async_request_cancel` + * (which becomes a no-op after completion), and `tikv_async_request_destroy`. + */ +typedef void (*tikv_async_completion_cb)(tikv_async_request_t *request, + tikv_status_t final_status, + void *userdata); + +/* + * Create a client. + * + * `pd_endpoints_csv` must be valid UTF-8 and contain at least one endpoint. + * `timeout_secs == 0` leaves the client timeout at its library default. + * `use_default_keyspace` enables API v2 default keyspace configuration. + * For compatibility, this constructor has no gRPC response decoding limit. + * Use the constructor below to choose a limit explicitly. + */ +tikv_status_t tikv_async_client_new(const char *pd_endpoints_csv, + uint32_t timeout_secs, + bool use_default_keyspace, + tikv_async_client_t **out_client, + tikv_error_t **out_error); + +/* + * Create a client with an explicit gRPC response decoding limit in bytes. + * + * `grpc_max_decoding_message_size == 0` uses the Rust client's default + * (currently 4 MiB). Larger responses require an explicit limit. + */ +tikv_status_t tikv_async_client_new_with_grpc_max_decoding_message_size( + const char *pd_endpoints_csv, + uint32_t timeout_secs, + bool use_default_keyspace, + size_t grpc_max_decoding_message_size, + tikv_async_client_t **out_client, + tikv_error_t **out_error); + +/* + * Start shutdown. + * + * This function is idempotent. + * + * After the first successful call, every later submit on the same client is + * rejected with `TIKV_STATUS_REJECTED`, and no callback is scheduled for those + * rejected submits. + * + * If `cancel_inflight` is true, every currently accepted in-flight request is + * asked to stop and will complete through its callback with + * `TIKV_STATUS_CANCELLED`. + */ +tikv_status_t tikv_async_client_shutdown_begin(tikv_async_client_t *client, + bool cancel_inflight, + tikv_error_t **out_error); + +/* + * Wait until every accepted in-flight request has finished its callback. + * + * This function may be called before or after shutdown begins. + * If shutdown has not yet begun, it only waits for the currently in-flight set; + * it does not prevent new submissions racing from other threads. + */ +tikv_status_t tikv_async_client_drain(const tikv_async_client_t *client, + tikv_error_t **out_error); + +/* + * Submit a raw put. + * + * Submit success means: + * + * - Rust copied the input buffers. + * - Rust registered the request for eventual completion. + * - Rust will invoke `completion` exactly once. + */ +tikv_status_t tikv_async_raw_put_submit(const tikv_async_client_t *client, + tikv_slice_t key, + tikv_slice_t value, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a raw put with a relative TTL in seconds. + * + * `ttl_secs == 0` disables TTL for the key. TiKV must be configured to support + * RawKV TTL before callers submit a non-zero value. + * + * Submit success has the same copy and completion guarantees as + * `tikv_async_raw_put_submit`. + */ +tikv_status_t tikv_async_raw_put_with_ttl_submit( + const tikv_async_client_t *client, + tikv_slice_t key, + tikv_slice_t value, + uint64_t ttl_secs, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a raw get. + * + * On success, completion with `TIKV_STATUS_OK` makes the value available + * through `tikv_async_request_take_get_result`. + */ +tikv_status_t tikv_async_raw_get_submit(const tikv_async_client_t *client, + tikv_slice_t key, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a request for the remaining RawKV TTL of a key. + * + * On success, completion with `TIKV_STATUS_OK` makes the result available + * through `tikv_async_request_take_get_key_ttl_result`. + */ +tikv_status_t tikv_async_raw_get_key_ttl_submit( + const tikv_async_client_t *client, + tikv_slice_t key, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a raw delete. + */ +tikv_status_t tikv_async_raw_delete_submit(const tikv_async_client_t *client, + tikv_slice_t key, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a raw delete-range. + * + * An empty `end_key` means unbounded upper range. + */ +tikv_status_t + tikv_async_raw_delete_range_submit(const tikv_async_client_t *client, + tikv_slice_t start_key, + tikv_slice_t end_key, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a forward raw scan. + * + * On success, completion with `TIKV_STATUS_OK` makes the pairs available + * through `tikv_async_request_take_scan_result`. + */ +tikv_status_t tikv_async_raw_scan_submit(const tikv_async_client_t *client, + tikv_slice_t start_key, + tikv_slice_t end_key, + uint32_t limit, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a backward raw scan. + * + * This maps to `client-rust` raw `scan_reverse`. + * The returned pairs are in reverse lexicographical key order. + * + * Note: the underlying raw client does not support reverse scan from an empty + * start key because locating the last region is not implemented. + */ +tikv_status_t + tikv_async_raw_scan_reverse_submit(const tikv_async_client_t *client, + tikv_slice_t start_key, + tikv_slice_t end_key, + uint32_t limit, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a mixed raw write batch. + * + * This reduces per-request overhead by letting the caller hand multiple put and + * delete mutations to one async request. + * + * `client-rust` raw API exposes `batch_put` and `batch_delete`, not one mixed + * mutation RPC. This binding therefore preserves the caller's put/delete phase + * order while coalescing adjacent same-type mutations into as few TiKV calls as + * practical. + * + * Example: + * + * PUT a, PUT b, DELETE c, DELETE d, PUT e + * + * becomes: + * + * batch_put(a, b) -> batch_delete(c, d) -> batch_put(e) + * + * Repeated PUTs to the same key inside one contiguous PUT run are split before + * submit so that later writes still win deterministically. + */ +tikv_status_t tikv_async_raw_batch_write_submit( + const tikv_async_client_t *client, + tikv_async_raw_mutations_t mutations, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a TTL-capable mixed raw write batch. + * + * This is the ABI-compatible TTL counterpart of + * `tikv_async_raw_batch_write_submit`. Each PUT may carry its own relative TTL + * in `ttl_secs`; PUTs with `ttl_secs == 0` keep the existing non-TTL behavior. + * DELETE ignores `ttl_secs`. + * + * The implementation maps PUT runs to `client-rust` + * `batch_put_with_ttl` and DELETE runs to `batch_delete`, preserving mutation + * phase order and later-write-wins behavior for repeated keys. + * + * TiKV must be configured to support RawKV TTL before callers submit a + * non-zero TTL. + */ +tikv_status_t tikv_async_raw_batch_write_with_ttl_submit( + const tikv_async_client_t *client, + tikv_async_raw_mutations_with_ttl_t mutations, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Submit a raw batch-get. + * + * TiKV raw batch-get returns only existing entries and does not preserve the + * input key order. + * On success, completion with `TIKV_STATUS_OK` makes the pairs available + * through `tikv_async_request_take_batch_get_result`. + */ +tikv_status_t tikv_async_raw_batch_get_submit(const tikv_async_client_t *client, + tikv_slices_t keys, + tikv_async_completion_cb completion, + void *userdata, + tikv_async_request_t **out_request, + tikv_error_t **out_error); + +/* + * Request lifecycle helpers. + */ + +/* + * Ask an accepted request to cancel. + * + * This function is best-effort and idempotent. + * + * - `*out_cancel_requested == true` means this call newly requested + * cancellation. + * - `*out_cancel_requested == false` means the request had already completed or + * cancellation had already been requested earlier. + * + * If the request was still pending, its completion callback still runs exactly + * once with `TIKV_STATUS_CANCELLED`. + */ +tikv_status_t tikv_async_request_cancel(tikv_async_request_t *request, + bool *out_cancel_requested, + tikv_error_t **out_error); + +/* + * Return whether the request has completed. + */ +tikv_status_t tikv_async_request_is_complete(const tikv_async_request_t *request, + bool *out_complete, + tikv_error_t **out_error); + +/* + * Return the terminal request status. + * + * This may be called only after completion. + */ +tikv_status_t tikv_async_request_completion_status( + const tikv_async_request_t *request, + tikv_status_t *out_status, + tikv_error_t **out_error); + +/* + * Return the terminal request error message, or NULL if none. + * + * The returned pointer is borrowed from the request state and remains valid + * until `tikv_async_request_destroy`. + */ +const char *tikv_async_request_error_message(const tikv_async_request_t *request); + +/* + * Result extraction helpers. + * + * Each successful result may be taken only once. + * They may be called only after completion, and only for a matching request + * kind that completed with `TIKV_STATUS_OK`. + */ +tikv_status_t tikv_async_request_take_get_result( + const tikv_async_request_t *request, + tikv_owned_bytes_t *out_value, + bool *out_found, + tikv_error_t **out_error); + +/* + * Take a get-key-TTL result. + * + * `*out_found == false` means the key does not exist and `*out_ttl_secs` is 0. + * `*out_found == true && *out_ttl_secs == 0` means the key exists without a + * TTL. Otherwise `*out_ttl_secs` is the remaining TTL in seconds. + */ +tikv_status_t tikv_async_request_take_get_key_ttl_result( + const tikv_async_request_t *request, + uint64_t *out_ttl_secs, + bool *out_found, + tikv_error_t **out_error); + +tikv_status_t tikv_async_request_take_scan_result( + const tikv_async_request_t *request, + tikv_kv_pairs_t *out_pairs, + tikv_error_t **out_error); + +tikv_status_t tikv_async_request_take_batch_get_result( + const tikv_async_request_t *request, + tikv_kv_pairs_t *out_pairs, + tikv_error_t **out_error); + +/* + * Destroy helpers. + */ +const char *tikv_error_message(const tikv_error_t *error); +void tikv_error_destroy(tikv_error_t *error); + +void tikv_owned_bytes_destroy(tikv_owned_bytes_t bytes); +void tikv_kv_pairs_destroy(tikv_kv_pairs_t pairs); + +void tikv_async_request_destroy(tikv_async_request_t *request); +void tikv_async_client_destroy(tikv_async_client_t *client); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/bindings/c/include/tikv_client_metrics.hpp b/bindings/c/include/tikv_client_metrics.hpp new file mode 100644 index 00000000..abe27865 --- /dev/null +++ b/bindings/c/include/tikv_client_metrics.hpp @@ -0,0 +1,620 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#ifndef TIKV_CLIENT_METRICS_HPP +#define TIKV_CLIENT_METRICS_HPP + +#include "tikv_client_async.hpp" +#include "tikv_client_metrics_c.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tikv { +namespace metrics { + +class MetricName final { +public: + constexpr explicit MetricName(std::string_view value) noexcept + : value_(value) {} + + constexpr std::string_view value() const noexcept { return value_; } + + friend constexpr bool operator==(MetricName lhs, MetricName rhs) noexcept { + return lhs.value_ == rhs.value_; + } + + friend constexpr bool operator!=(MetricName lhs, MetricName rhs) noexcept { + return !(lhs == rhs); + } + +private: + std::string_view value_; +}; + +class MetricTypeLabel final { +public: + constexpr explicit MetricTypeLabel(std::string_view value) noexcept + : value_(value) {} + + constexpr std::string_view value() const noexcept { return value_; } + + friend constexpr bool operator==(MetricTypeLabel lhs, + MetricTypeLabel rhs) noexcept { + return lhs.value_ == rhs.value_; + } + + friend constexpr bool operator!=(MetricTypeLabel lhs, + MetricTypeLabel rhs) noexcept { + return !(lhs == rhs); + } + +private: + std::string_view value_; +}; + +namespace metric_names { + +inline constexpr MetricName TikvRequestTotal{ + TIKV_METRIC_TIKV_REQUEST_TOTAL}; +inline constexpr MetricName TikvRequestDurationSeconds{ + TIKV_METRIC_TIKV_REQUEST_DURATION_SECONDS}; +inline constexpr MetricName TikvFailedRequestTotal{ + TIKV_METRIC_TIKV_FAILED_REQUEST_TOTAL}; +inline constexpr MetricName TikvFailedRequestDurationSeconds{ + TIKV_METRIC_TIKV_FAILED_REQUEST_DURATION_SECONDS}; +inline constexpr MetricName PdRequestTotal{TIKV_METRIC_PD_REQUEST_TOTAL}; +inline constexpr MetricName PdRequestDurationSeconds{ + TIKV_METRIC_PD_REQUEST_DURATION_SECONDS}; +inline constexpr MetricName PdFailedRequestTotal{ + TIKV_METRIC_PD_FAILED_REQUEST_TOTAL}; +inline constexpr MetricName PdFailedRequestDurationSeconds{ + TIKV_METRIC_PD_FAILED_REQUEST_DURATION_SECONDS}; +inline constexpr MetricName PdTsoBatchSize{TIKV_METRIC_PD_TSO_BATCH_SIZE}; + +} // namespace metric_names + +namespace metric_type_labels { + +inline constexpr MetricTypeLabel RawGet{TIKV_METRIC_LABEL_TYPE_RAW_GET}; +inline constexpr MetricTypeLabel RawBatchGet{ + TIKV_METRIC_LABEL_TYPE_RAW_BATCH_GET}; +inline constexpr MetricTypeLabel RawGetKeyTtl{ + TIKV_METRIC_LABEL_TYPE_RAW_GET_KEY_TTL}; +inline constexpr MetricTypeLabel RawPut{TIKV_METRIC_LABEL_TYPE_RAW_PUT}; +inline constexpr MetricTypeLabel RawBatchPut{ + TIKV_METRIC_LABEL_TYPE_RAW_BATCH_PUT}; +inline constexpr MetricTypeLabel RawDelete{ + TIKV_METRIC_LABEL_TYPE_RAW_DELETE}; +inline constexpr MetricTypeLabel RawBatchDelete{ + TIKV_METRIC_LABEL_TYPE_RAW_BATCH_DELETE}; +inline constexpr MetricTypeLabel RawScan{TIKV_METRIC_LABEL_TYPE_RAW_SCAN}; +inline constexpr MetricTypeLabel RawBatchScan{ + TIKV_METRIC_LABEL_TYPE_RAW_BATCH_SCAN}; +inline constexpr MetricTypeLabel RawDeleteRange{ + TIKV_METRIC_LABEL_TYPE_RAW_DELETE_RANGE}; +inline constexpr MetricTypeLabel RawCompareAndSwap{ + TIKV_METRIC_LABEL_TYPE_RAW_COMPARE_AND_SWAP}; +inline constexpr MetricTypeLabel RawCoprocessor{ + TIKV_METRIC_LABEL_TYPE_RAW_COPROCESSOR}; +inline constexpr MetricTypeLabel KvGet{TIKV_METRIC_LABEL_TYPE_KV_GET}; +inline constexpr MetricTypeLabel KvScan{TIKV_METRIC_LABEL_TYPE_KV_SCAN}; +inline constexpr MetricTypeLabel KvPrewrite{ + TIKV_METRIC_LABEL_TYPE_KV_PREWRITE}; +inline constexpr MetricTypeLabel KvCommit{TIKV_METRIC_LABEL_TYPE_KV_COMMIT}; +inline constexpr MetricTypeLabel KvBatchGet{ + TIKV_METRIC_LABEL_TYPE_KV_BATCH_GET}; +inline constexpr MetricTypeLabel KvBatchRollback{ + TIKV_METRIC_LABEL_TYPE_KV_BATCH_ROLLBACK}; +inline constexpr MetricTypeLabel KvPessimisticRollback{ + TIKV_METRIC_LABEL_TYPE_KV_PESSIMISTIC_ROLLBACK}; +inline constexpr MetricTypeLabel KvResolveLock{ + TIKV_METRIC_LABEL_TYPE_KV_RESOLVE_LOCK}; +inline constexpr MetricTypeLabel KvScanLock{ + TIKV_METRIC_LABEL_TYPE_KV_SCAN_LOCK}; +inline constexpr MetricTypeLabel KvPessimisticLock{ + TIKV_METRIC_LABEL_TYPE_KV_PESSIMISTIC_LOCK}; +inline constexpr MetricTypeLabel KvTxnHeartBeat{ + TIKV_METRIC_LABEL_TYPE_KV_TXN_HEART_BEAT}; +inline constexpr MetricTypeLabel KvCheckTxnStatus{ + TIKV_METRIC_LABEL_TYPE_KV_CHECK_TXN_STATUS}; +inline constexpr MetricTypeLabel KvCheckSecondaryLocksRequest{ + TIKV_METRIC_LABEL_TYPE_KV_CHECK_SECONDARY_LOCKS_REQUEST}; +inline constexpr MetricTypeLabel KvGc{TIKV_METRIC_LABEL_TYPE_KV_GC}; +inline constexpr MetricTypeLabel KvDeleteRange{ + TIKV_METRIC_LABEL_TYPE_KV_DELETE_RANGE}; +inline constexpr MetricTypeLabel UnsafeDestroyRange{ + TIKV_METRIC_LABEL_TYPE_UNSAFE_DESTROY_RANGE}; +inline constexpr MetricTypeLabel PdGetRegion{ + TIKV_METRIC_LABEL_TYPE_PD_GET_REGION}; +inline constexpr MetricTypeLabel PdGetRegionById{ + TIKV_METRIC_LABEL_TYPE_PD_GET_REGION_BY_ID}; +inline constexpr MetricTypeLabel PdGetStore{ + TIKV_METRIC_LABEL_TYPE_PD_GET_STORE}; +inline constexpr MetricTypeLabel PdGetAllStores{ + TIKV_METRIC_LABEL_TYPE_PD_GET_ALL_STORES}; +inline constexpr MetricTypeLabel PdGetTimestamp{ + TIKV_METRIC_LABEL_TYPE_PD_GET_TIMESTAMP}; +inline constexpr MetricTypeLabel PdUpdateGcSafepoint{ + TIKV_METRIC_LABEL_TYPE_PD_UPDATE_GC_SAFEPOINT}; +inline constexpr MetricTypeLabel PdLoadKeyspace{ + TIKV_METRIC_LABEL_TYPE_PD_LOAD_KEYSPACE}; + +} // namespace metric_type_labels + +using MetricSampleKind = ::tikv_metric_sample_kind_t; + +namespace metric_sample_kinds { + +inline constexpr MetricSampleKind Unknown = TIKV_METRIC_SAMPLE_UNKNOWN; +inline constexpr MetricSampleKind Counter = TIKV_METRIC_SAMPLE_COUNTER; +inline constexpr MetricSampleKind HistogramBucket = + TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET; +inline constexpr MetricSampleKind HistogramSum = + TIKV_METRIC_SAMPLE_HISTOGRAM_SUM; +inline constexpr MetricSampleKind HistogramCount = + TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT; + +} // namespace metric_sample_kinds + +struct MetricSample { + std::string name; + std::optional type_label; + MetricSampleKind kind = metric_sample_kinds::Unknown; + double value = 0.0; + double le = 0.0; + + static MetricSample from_c(const tikv_metric_sample_t &sample) { + MetricSample result; + if (sample.name != nullptr) { + result.name = sample.name; + } + if (sample.type_label != nullptr) { + result.type_label = sample.type_label; + } + result.kind = sample.kind; + result.value = sample.value; + result.le = sample.le; + return result; + } + + const std::string &name_string() const noexcept { return name; } + + const std::optional &type_label_string() const noexcept { + return type_label; + } + + bool kind_is(MetricSampleKind expected) const noexcept { + return kind == expected; + } + + bool name_is(MetricName expected) const noexcept { + return std::string_view(name.data(), name.size()) == expected.value(); + } + + bool has_no_type_label() const noexcept { + return !type_label.has_value(); + } + + bool type_label_is(MetricTypeLabel expected) const noexcept { + return type_label.has_value() && + std::string_view(type_label->data(), type_label->size()) == + expected.value(); + } + + bool type_label_is(std::optional expected) const noexcept { + if (!expected.has_value()) { + return has_no_type_label(); + } + return type_label_is(*expected); + } + + bool is_counter(MetricName expected_name) const noexcept { + return kind_is(metric_sample_kinds::Counter) && name_is(expected_name); + } + + bool is_counter(MetricName expected_name, + MetricTypeLabel expected_label) const noexcept { + return is_counter(expected_name) && type_label_is(expected_label); + } + + bool is_counter(MetricName expected_name, + std::optional expected_label) const noexcept { + return is_counter(expected_name) && type_label_is(expected_label); + } + + bool is_histogram_bucket(MetricName expected_name) const noexcept { + return kind_is(metric_sample_kinds::HistogramBucket) && + name_is(expected_name); + } + + bool is_histogram_bucket(MetricName expected_name, + MetricTypeLabel expected_label) const noexcept { + return is_histogram_bucket(expected_name) && type_label_is(expected_label); + } + + bool is_histogram_bucket( + MetricName expected_name, + std::optional expected_label) const noexcept { + return is_histogram_bucket(expected_name) && type_label_is(expected_label); + } + + bool is_histogram_sum(MetricName expected_name) const noexcept { + return kind_is(metric_sample_kinds::HistogramSum) && + name_is(expected_name); + } + + bool is_histogram_sum(MetricName expected_name, + MetricTypeLabel expected_label) const noexcept { + return is_histogram_sum(expected_name) && type_label_is(expected_label); + } + + bool is_histogram_sum( + MetricName expected_name, + std::optional expected_label) const noexcept { + return is_histogram_sum(expected_name) && type_label_is(expected_label); + } + + bool is_histogram_count(MetricName expected_name) const noexcept { + return kind_is(metric_sample_kinds::HistogramCount) && + name_is(expected_name); + } + + bool is_histogram_count(MetricName expected_name, + MetricTypeLabel expected_label) const noexcept { + return is_histogram_count(expected_name) && type_label_is(expected_label); + } + + bool is_histogram_count( + MetricName expected_name, + std::optional expected_label) const noexcept { + return is_histogram_count(expected_name) && type_label_is(expected_label); + } +}; + +class MetricSampleView : public ::tikv::async::BorrowedResultValueTag { +public: + explicit MetricSampleView(const tikv_metric_sample_t &sample) noexcept + : sample_(sample) {} + + const tikv_metric_sample_t &raw() const noexcept { return sample_; } + + std::string_view name() const noexcept { + if (sample_.name == nullptr) { + return {}; + } + return std::string_view(sample_.name); + } + + std::optional type_label() const noexcept { + if (sample_.type_label == nullptr) { + return std::nullopt; + } + return std::string_view(sample_.type_label); + } + + MetricSampleKind kind() const noexcept { + return sample_.kind; + } + + bool kind_is(MetricSampleKind kind) const noexcept { + return sample_.kind == kind; + } + + bool name_is(MetricName name) const noexcept { + return sample_.name != nullptr && + std::string_view(sample_.name) == name.value(); + } + + bool has_no_type_label() const noexcept { + return sample_.type_label == nullptr; + } + + bool type_label_is(MetricTypeLabel label) const noexcept { + return type_label_is(std::optional(label)); + } + + bool type_label_is(std::optional label) const noexcept { + if (!label.has_value()) { + return has_no_type_label(); + } + return sample_.type_label != nullptr && + std::string_view(sample_.type_label) == label->value(); + } + + bool is_counter(MetricName name) const noexcept { + return kind_is(metric_sample_kinds::Counter) && name_is(name); + } + + bool is_counter(MetricName name, MetricTypeLabel label) const noexcept { + return is_counter(name) && type_label_is(label); + } + + bool is_counter(MetricName name, + std::optional label) const noexcept { + return is_counter(name) && type_label_is(label); + } + + bool is_histogram_bucket(MetricName name) const noexcept { + return kind_is(metric_sample_kinds::HistogramBucket) && name_is(name); + } + + bool is_histogram_bucket(MetricName name, + MetricTypeLabel label) const noexcept { + return is_histogram_bucket(name) && type_label_is(label); + } + + bool is_histogram_bucket(MetricName name, + std::optional label) const noexcept { + return is_histogram_bucket(name) && type_label_is(label); + } + + bool is_histogram_sum(MetricName name) const noexcept { + return kind_is(metric_sample_kinds::HistogramSum) && name_is(name); + } + + bool is_histogram_sum(MetricName name, MetricTypeLabel label) const noexcept { + return is_histogram_sum(name) && type_label_is(label); + } + + bool is_histogram_sum(MetricName name, + std::optional label) const noexcept { + return is_histogram_sum(name) && type_label_is(label); + } + + bool is_histogram_count(MetricName name) const noexcept { + return kind_is(metric_sample_kinds::HistogramCount) && name_is(name); + } + + bool is_histogram_count(MetricName name, MetricTypeLabel label) const noexcept { + return is_histogram_count(name) && type_label_is(label); + } + + bool is_histogram_count(MetricName name, + std::optional label) const noexcept { + return is_histogram_count(name) && type_label_is(label); + } + + double value() const noexcept { return sample_.value; } + + double le() const noexcept { return sample_.le; } + + MetricSample to_owned() const { return MetricSample::from_c(sample_); } + +private: + tikv_metric_sample_t sample_; +}; + +class MetricsSnapshot { +public: + using Result = ::tikv::async::Result; + using SampleResult = ::tikv::async::Result; + using SampleViewResult = ::tikv::async::Result; + using SamplesResult = ::tikv::async::Result>; + + class const_iterator { + public: + using iterator_category = std::input_iterator_tag; + using value_type = SampleResult; + using difference_type = std::ptrdiff_t; + using pointer = void; + using reference = value_type; + + const_iterator() noexcept = default; + + const_iterator(const MetricsSnapshot *snapshot, std::size_t index) noexcept + : snapshot_(snapshot), index_(index) {} + + value_type operator*() const { + if (snapshot_ == nullptr) { + return ::tikv::async::Status::InvalidArgument( + "metrics snapshot iterator is empty"); + } + return snapshot_->get(index_); + } + + const_iterator &operator++() noexcept { + ++index_; + return *this; + } + + const_iterator operator++(int) noexcept { + const_iterator current(*this); + ++(*this); + return current; + } + + bool operator==(const const_iterator &other) const noexcept { + return snapshot_ == other.snapshot_ && index_ == other.index_; + } + + bool operator!=(const const_iterator &other) const noexcept { + return !(*this == other); + } + + private: + const MetricsSnapshot *snapshot_ = nullptr; + std::size_t index_ = 0; + }; + + class view_iterator { + public: + using iterator_category = std::input_iterator_tag; + using value_type = SampleViewResult; + using difference_type = std::ptrdiff_t; + using pointer = void; + using reference = value_type; + + view_iterator() noexcept = default; + + view_iterator(const MetricsSnapshot *snapshot, std::size_t index) noexcept + : snapshot_(snapshot), index_(index) {} + + value_type operator*() const { + if (snapshot_ == nullptr) { + return ::tikv::async::Status::InvalidArgument( + "metrics snapshot view iterator is empty"); + } + return snapshot_->get_view(index_); + } + + view_iterator &operator++() noexcept { + ++index_; + return *this; + } + + view_iterator operator++(int) noexcept { + view_iterator current(*this); + ++(*this); + return current; + } + + bool operator==(const view_iterator &other) const noexcept { + return snapshot_ == other.snapshot_ && index_ == other.index_; + } + + bool operator!=(const view_iterator &other) const noexcept { + return !(*this == other); + } + + private: + const MetricsSnapshot *snapshot_ = nullptr; + std::size_t index_ = 0; + }; + + class ViewRange { + public: + explicit ViewRange(const MetricsSnapshot *snapshot) noexcept + : snapshot_(snapshot) {} + + view_iterator begin() const noexcept { + return view_iterator(snapshot_, 0); + } + + view_iterator end() const noexcept { + return view_iterator(snapshot_, + snapshot_ == nullptr ? 0 : snapshot_->size()); + } + + private: + const MetricsSnapshot *snapshot_ = nullptr; + }; + + MetricsSnapshot() noexcept = default; + + MetricsSnapshot(const MetricsSnapshot &) = delete; + MetricsSnapshot &operator=(const MetricsSnapshot &) = delete; + + MetricsSnapshot(MetricsSnapshot &&other) noexcept : raw_(other.release()) {} + + MetricsSnapshot &operator=(MetricsSnapshot &&other) noexcept { + if (this != &other) { + reset(other.release()); + } + return *this; + } + + ~MetricsSnapshot() { reset(); } + + static Result collect() { + tikv_metrics_snapshot_t *raw = nullptr; + tikv_error_t *error = nullptr; + const tikv_status_t code = tikv_metrics_snapshot_collect(&raw, &error); + ::tikv::async::Status status = ::tikv::async::Status::from_ffi(code, error); + if (!status.ok()) { + return status; + } + return MetricsSnapshot(raw); + } + + bool valid() const noexcept { return raw_ != nullptr; } + + explicit operator bool() const noexcept { return valid(); } + + std::size_t size() const noexcept { + return tikv_metrics_snapshot_len(raw_); + } + + const_iterator begin() const noexcept { return const_iterator(this, 0); } + + const_iterator end() const noexcept { return const_iterator(this, size()); } + + view_iterator view_begin() const noexcept { return view_iterator(this, 0); } + + view_iterator view_end() const noexcept { + return view_iterator(this, size()); + } + + ViewRange views() const noexcept { return ViewRange(this); } + + SampleResult get(std::size_t index) const { + auto raw_sample = get_raw_sample(index); + if (!raw_sample.ok()) { + return raw_sample.status(); + } + return MetricSample::from_c(raw_sample.value()); + } + + SampleViewResult get_view(std::size_t index) const { + auto raw_sample = get_raw_sample(index); + if (!raw_sample.ok()) { + return raw_sample.status(); + } + return MetricSampleView(raw_sample.value()); + } + + SamplesResult samples() const { + std::vector result; + result.reserve(size()); + for (std::size_t i = 0; i < size(); ++i) { + auto sample = get(i); + if (!sample.ok()) { + return sample.status(); + } + result.push_back(sample.take()); + } + return result; + } + + tikv_metrics_snapshot_t *get_raw() const noexcept { return raw_; } + + tikv_metrics_snapshot_t *release() noexcept { + tikv_metrics_snapshot_t *raw = raw_; + raw_ = nullptr; + return raw; + } + + void reset(tikv_metrics_snapshot_t *raw = nullptr) noexcept { + if (raw_ != nullptr) { + tikv_metrics_snapshot_destroy(raw_); + } + raw_ = raw; + } + +private: + explicit MetricsSnapshot(tikv_metrics_snapshot_t *raw) noexcept : raw_(raw) {} + + ::tikv::async::Result + get_raw_sample(std::size_t index) const { + tikv_metric_sample_t sample = {}; + tikv_error_t *error = nullptr; + const tikv_status_t code = + tikv_metrics_snapshot_get(raw_, index, &sample, &error); + ::tikv::async::Status status = ::tikv::async::Status::from_ffi(code, error); + if (!status.ok()) { + return status; + } + return sample; + } + + tikv_metrics_snapshot_t *raw_ = nullptr; +}; + +} // namespace metrics +} // namespace tikv + +#endif // TIKV_CLIENT_METRICS_HPP diff --git a/bindings/c/include/tikv_client_metrics_c.h b/bindings/c/include/tikv_client_metrics_c.h new file mode 100644 index 00000000..5b82e446 --- /dev/null +++ b/bindings/c/include/tikv_client_metrics_c.h @@ -0,0 +1,116 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#ifndef TIKV_CLIENT_METRICS_C_H +#define TIKV_CLIENT_METRICS_C_H + +#include "tikv_client_async_c.h" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct tikv_metrics_snapshot_t tikv_metrics_snapshot_t; + +typedef enum tikv_metric_sample_kind_t { + TIKV_METRIC_SAMPLE_UNKNOWN = 0, + TIKV_METRIC_SAMPLE_COUNTER = 1, + TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET = 2, + TIKV_METRIC_SAMPLE_HISTOGRAM_SUM = 3, + TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT = 4 +} tikv_metric_sample_kind_t; + +/* + * Metric names currently emitted by client-rust. + */ +#define TIKV_METRIC_TIKV_REQUEST_TOTAL "tikv_request_total" +#define TIKV_METRIC_TIKV_REQUEST_DURATION_SECONDS \ + "tikv_request_duration_seconds" +#define TIKV_METRIC_TIKV_FAILED_REQUEST_TOTAL "tikv_failed_request_total" +#define TIKV_METRIC_TIKV_FAILED_REQUEST_DURATION_SECONDS \ + "tikv_failed_request_duration_seconds" +#define TIKV_METRIC_PD_REQUEST_TOTAL "pd_request_total" +#define TIKV_METRIC_PD_REQUEST_DURATION_SECONDS "pd_request_duration_seconds" +#define TIKV_METRIC_PD_FAILED_REQUEST_TOTAL "pd_failed_request_total" +#define TIKV_METRIC_PD_FAILED_REQUEST_DURATION_SECONDS \ + "pd_failed_request_duration_seconds" +#define TIKV_METRIC_PD_TSO_BATCH_SIZE "pd_tso_batch_size" + +/* + * Values currently used for the Prometheus label key `type`. + */ +#define TIKV_METRIC_LABEL_TYPE_RAW_GET "raw_get" +#define TIKV_METRIC_LABEL_TYPE_RAW_BATCH_GET "raw_batch_get" +#define TIKV_METRIC_LABEL_TYPE_RAW_GET_KEY_TTL "raw_get_key_ttl" +#define TIKV_METRIC_LABEL_TYPE_RAW_PUT "raw_put" +#define TIKV_METRIC_LABEL_TYPE_RAW_BATCH_PUT "raw_batch_put" +#define TIKV_METRIC_LABEL_TYPE_RAW_DELETE "raw_delete" +#define TIKV_METRIC_LABEL_TYPE_RAW_BATCH_DELETE "raw_batch_delete" +#define TIKV_METRIC_LABEL_TYPE_RAW_SCAN "raw_scan" +#define TIKV_METRIC_LABEL_TYPE_RAW_BATCH_SCAN "raw_batch_scan" +#define TIKV_METRIC_LABEL_TYPE_RAW_DELETE_RANGE "raw_delete_range" +#define TIKV_METRIC_LABEL_TYPE_RAW_COMPARE_AND_SWAP "raw_compare_and_swap" +#define TIKV_METRIC_LABEL_TYPE_RAW_COPROCESSOR "raw_coprocessor" + +#define TIKV_METRIC_LABEL_TYPE_KV_GET "kv_get" +#define TIKV_METRIC_LABEL_TYPE_KV_SCAN "kv_scan" +#define TIKV_METRIC_LABEL_TYPE_KV_PREWRITE "kv_prewrite" +#define TIKV_METRIC_LABEL_TYPE_KV_COMMIT "kv_commit" +#define TIKV_METRIC_LABEL_TYPE_KV_BATCH_GET "kv_batch_get" +#define TIKV_METRIC_LABEL_TYPE_KV_BATCH_ROLLBACK "kv_batch_rollback" +#define TIKV_METRIC_LABEL_TYPE_KV_PESSIMISTIC_ROLLBACK \ + "kv_pessimistic_rollback" +#define TIKV_METRIC_LABEL_TYPE_KV_RESOLVE_LOCK "kv_resolve_lock" +#define TIKV_METRIC_LABEL_TYPE_KV_SCAN_LOCK "kv_scan_lock" +#define TIKV_METRIC_LABEL_TYPE_KV_PESSIMISTIC_LOCK "kv_pessimistic_lock" +#define TIKV_METRIC_LABEL_TYPE_KV_TXN_HEART_BEAT "kv_txn_heart_beat" +#define TIKV_METRIC_LABEL_TYPE_KV_CHECK_TXN_STATUS "kv_check_txn_status" +#define TIKV_METRIC_LABEL_TYPE_KV_CHECK_SECONDARY_LOCKS_REQUEST \ + "kv_check_secondary_locks_request" +#define TIKV_METRIC_LABEL_TYPE_KV_GC "kv_gc" +#define TIKV_METRIC_LABEL_TYPE_KV_DELETE_RANGE "kv_delete_range" +#define TIKV_METRIC_LABEL_TYPE_UNSAFE_DESTROY_RANGE "unsafe_destroy_range" + +#define TIKV_METRIC_LABEL_TYPE_PD_GET_REGION "get_region" +#define TIKV_METRIC_LABEL_TYPE_PD_GET_REGION_BY_ID "get_region_by_id" +#define TIKV_METRIC_LABEL_TYPE_PD_GET_STORE "get_store" +#define TIKV_METRIC_LABEL_TYPE_PD_GET_ALL_STORES "get_all_stores" +#define TIKV_METRIC_LABEL_TYPE_PD_GET_TIMESTAMP "get_timestamp" +#define TIKV_METRIC_LABEL_TYPE_PD_UPDATE_GC_SAFEPOINT \ + "update_gc_safepoint" +#define TIKV_METRIC_LABEL_TYPE_PD_LOAD_KEYSPACE "load_keyspace" + +typedef struct tikv_metric_sample_t { + const char *name; + const char *type_label; + tikv_metric_sample_kind_t kind; + double value; + double le; +} tikv_metric_sample_t; + +/* + * Collect a structured snapshot of the current process Prometheus default + * registry. + * + * Strings referenced by `tikv_metric_sample_t` are owned by the snapshot and + * remain valid until `tikv_metrics_snapshot_destroy`. `type_label` is NULL when + * the sample has no Prometheus `type` label. + */ +tikv_status_t tikv_metrics_snapshot_collect( + tikv_metrics_snapshot_t **out_snapshot, + tikv_error_t **out_error); + +size_t tikv_metrics_snapshot_len(const tikv_metrics_snapshot_t *snapshot); + +tikv_status_t tikv_metrics_snapshot_get( + const tikv_metrics_snapshot_t *snapshot, + size_t index, + tikv_metric_sample_t *out_sample, + tikv_error_t **out_error); + +void tikv_metrics_snapshot_destroy(tikv_metrics_snapshot_t *snapshot); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/bindings/c/src/lib.rs b/bindings/c/src/lib.rs new file mode 100644 index 00000000..afb7318a --- /dev/null +++ b/bindings/c/src/lib.rs @@ -0,0 +1,2475 @@ +// Copyright (c) 2026 PingCAP, Inc. + +#![allow(non_camel_case_types)] + +use std::collections::HashSet; +use std::ffi::{c_char, c_int, c_void, CStr, CString}; +use std::mem; +use std::panic::{self, AssertUnwindSafe}; +use std::ptr; +use std::slice; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex, Weak}; +use std::time::Duration; + +use futures::FutureExt; +use prometheus::proto::MetricType; +use tikv_client::{BoundRange, Config, Key, KvPair, RawClient}; +use tokio::runtime::Handle; + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_slice_t { + pub ptr: *const u8, + pub len: usize, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_slices_t { + pub ptr: *const tikv_slice_t, + pub len: usize, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_async_raw_mutation_t { + pub op: tikv_async_raw_mutation_op_t, + pub key: tikv_slice_t, + pub value: tikv_slice_t, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_async_raw_mutations_t { + pub ptr: *const tikv_async_raw_mutation_t, + pub len: usize, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_async_raw_mutation_with_ttl_t { + pub op: tikv_async_raw_mutation_op_t, + pub key: tikv_slice_t, + pub value: tikv_slice_t, + pub ttl_secs: u64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_async_raw_mutations_with_ttl_t { + pub ptr: *const tikv_async_raw_mutation_with_ttl_t, + pub len: usize, +} + +pub type tikv_async_raw_mutation_op_t = c_int; + +const TIKV_ASYNC_RAW_MUTATION_OP_PUT: tikv_async_raw_mutation_op_t = 0; +const TIKV_ASYNC_RAW_MUTATION_OP_DELETE: tikv_async_raw_mutation_op_t = 1; + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_owned_bytes_t { + pub ptr: *mut u8, + pub len: usize, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_kv_pair_t { + pub key: tikv_owned_bytes_t, + pub value: tikv_owned_bytes_t, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_kv_pairs_t { + pub ptr: *mut tikv_kv_pair_t, + pub len: usize, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum tikv_status_t { + TIKV_STATUS_OK = 0, + TIKV_STATUS_ERROR = 1, + TIKV_STATUS_INVALID_ARGUMENT = 2, + TIKV_STATUS_PANIC = 3, + TIKV_STATUS_REJECTED = 4, + TIKV_STATUS_CANCELLED = 5, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum tikv_metric_sample_kind_t { + TIKV_METRIC_SAMPLE_UNKNOWN = 0, + TIKV_METRIC_SAMPLE_COUNTER = 1, + TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET = 2, + TIKV_METRIC_SAMPLE_HISTOGRAM_SUM = 3, + TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT = 4, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_metric_sample_t { + pub name: *const c_char, + pub type_label: *const c_char, + pub kind: tikv_metric_sample_kind_t, + pub value: f64, + pub le: f64, +} + +pub type tikv_async_completion_cb = extern "C" fn( + request: *mut tikv_async_request_t, + final_status: tikv_status_t, + userdata: *mut c_void, +); + +pub struct tikv_error_t { + message: CString, +} + +#[repr(C)] +pub struct tikv_async_client_t { + _private: [u8; 0], +} + +#[repr(C)] +pub struct tikv_async_request_t { + _private: [u8; 0], +} + +#[repr(C)] +pub struct tikv_metrics_snapshot_t { + _private: [u8; 0], +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RequestKind { + Put, + PutWithTtl, + Get, + GetKeyTtl, + Delete, + DeleteRange, + Scan, + ScanReverse, + BatchWrite, + BatchGet, +} + +impl RequestKind { + fn name(self) -> &'static str { + match self { + Self::Put => "put", + Self::PutWithTtl => "put_with_ttl", + Self::Get => "get", + Self::GetKeyTtl => "get_key_ttl", + Self::Delete => "delete", + Self::DeleteRange => "delete_range", + Self::Scan => "scan", + Self::ScanReverse => "scan_reverse", + Self::BatchWrite => "batch_write", + Self::BatchGet => "batch_get", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum MutationOp { + Put, + Delete, +} + +impl MutationOp { + fn from_ffi(raw: tikv_async_raw_mutation_op_t, name: &str) -> FfiResult { + match raw { + TIKV_ASYNC_RAW_MUTATION_OP_PUT => Ok(Self::Put), + TIKV_ASYNC_RAW_MUTATION_OP_DELETE => Ok(Self::Delete), + _ => Err(FfiError::invalid_argument(format!( + "{name}.op has invalid value {raw}" + ))), + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum OwnedMutation { + Put(KvPair, u64), + Delete(Key), +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum WriteBatchRun { + Put(Vec<(KvPair, u64)>), + Delete(Vec), +} + +#[derive(Debug)] +enum RequestOutput { + Unit, + Get(Option>), + KeyTtl(Option), + Pairs(Vec), +} + +#[derive(Clone, Debug)] +struct FfiError { + status: tikv_status_t, + message: String, +} + +type FfiResult = Result; +type PinBoxRequestFuture = + std::pin::Pin> + Send + 'static>>; + +impl FfiError { + fn error(message: impl Into) -> Self { + Self { + status: tikv_status_t::TIKV_STATUS_ERROR, + message: message.into(), + } + } + + fn invalid_argument(message: impl Into) -> Self { + Self { + status: tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT, + message: message.into(), + } + } + + fn rejected(message: impl Into) -> Self { + Self { + status: tikv_status_t::TIKV_STATUS_REJECTED, + message: message.into(), + } + } +} + +struct DrainState { + inflight: usize, +} + +struct ClientLifecycle { + accepting: AtomicBool, + cancel_inflight: AtomicBool, + drain: Mutex, + drain_ready: Condvar, + requests: Mutex>, +} + +struct AsyncClientCore { + runtime: tokio::runtime::Runtime, + client: RawClient, + lifecycle: ClientLifecycle, +} + +struct AsyncClientHandle { + core: Arc, +} + +struct CompletionState { + complete: bool, + callback_running: bool, + callback_done: bool, + completion_status: Option, + error_message: Option, + output: Option, + output_taken: bool, + cancel_requested: bool, +} + +struct RequestState { + kind: RequestKind, + completion: Mutex, + ready: Condvar, + client: Weak, + callback: Option, + userdata: *mut c_void, + retain_count: AtomicUsize, +} + +unsafe impl Send for RequestState {} +unsafe impl Sync for RequestState {} + +struct MetricsSnapshot { + samples: Vec, +} + +struct MetricSample { + ffi: tikv_metric_sample_t, + _name: CString, + _type_label: Option, +} + +fn ffi_cstring(value: &str) -> CString { + CString::new(value.replace('\0', " ")).expect("sanitized string") +} + +fn metric_type_label(metric: &prometheus::proto::Metric) -> Option<&str> { + metric + .get_label() + .iter() + .find(|label| label.get_name() == "type") + .map(|label| label.get_value()) +} + +fn new_metric_sample( + name: &str, + kind: tikv_metric_sample_kind_t, + type_label: Option<&str>, + value: f64, + le: f64, +) -> MetricSample { + let name = ffi_cstring(name); + let type_label = type_label.map(ffi_cstring); + let ffi = tikv_metric_sample_t { + name: name.as_ptr(), + type_label: type_label + .as_ref() + .map_or(ptr::null(), |label| label.as_ptr()), + kind, + value, + le, + }; + MetricSample { + ffi, + _name: name, + _type_label: type_label, + } +} + +fn collect_metrics_snapshot() -> MetricsSnapshot { + let mut samples = Vec::new(); + + for family in prometheus::gather() { + let name = family.get_name(); + + match family.get_field_type() { + MetricType::COUNTER => { + for metric in family.get_metric() { + let type_label = metric_type_label(metric); + samples.push(new_metric_sample( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_COUNTER, + type_label, + metric.get_counter().get_value(), + 0.0, + )); + } + } + MetricType::HISTOGRAM => { + for metric in family.get_metric() { + let type_label = metric_type_label(metric); + let histogram = metric.get_histogram(); + let mut inf_seen = false; + for bucket in histogram.get_bucket() { + let upper_bound = bucket.get_upper_bound(); + inf_seen |= upper_bound.is_sign_positive() && upper_bound.is_infinite(); + samples.push(new_metric_sample( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET, + type_label, + bucket.get_cumulative_count() as f64, + upper_bound, + )); + } + if !inf_seen { + samples.push(new_metric_sample( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET, + type_label, + histogram.get_sample_count() as f64, + f64::INFINITY, + )); + } + samples.push(new_metric_sample( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_SUM, + type_label, + histogram.get_sample_sum(), + 0.0, + )); + samples.push(new_metric_sample( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT, + type_label, + histogram.get_sample_count() as f64, + 0.0, + )); + } + } + _ => {} + } + } + + MetricsSnapshot { samples } +} + +fn metrics_snapshot_ref<'a>(raw: *const tikv_metrics_snapshot_t) -> FfiResult<&'a MetricsSnapshot> { + if raw.is_null() { + return Err(FfiError::invalid_argument("snapshot must not be null")); + } + Ok(unsafe { &*(raw as *const MetricsSnapshot) }) +} + +fn clear_error_slot(out_error: *mut *mut tikv_error_t) { + if out_error.is_null() { + return; + } + + unsafe { + *out_error = ptr::null_mut(); + } +} + +fn set_error_slot(out_error: *mut *mut tikv_error_t, message: String) { + if out_error.is_null() { + return; + } + + let sanitized = message.replace('\0', " "); + let c_message = CString::new(sanitized).unwrap_or_else(|_| { + CString::new("failed to build FFI error message").expect("static string") + }); + let error = Box::new(tikv_error_t { message: c_message }); + unsafe { + *out_error = Box::into_raw(error); + } +} + +fn panic_message(payload: Box) -> String { + if let Some(message) = payload.downcast_ref::<&str>() { + return format!("panic across FFI boundary: {message}"); + } + if let Some(message) = payload.downcast_ref::() { + return format!("panic across FFI boundary: {message}"); + } + "panic across FFI boundary".to_owned() +} + +fn run_ffi(out_error: *mut *mut tikv_error_t, f: F) -> Result +where + F: FnOnce() -> FfiResult, +{ + clear_error_slot(out_error); + + match panic::catch_unwind(AssertUnwindSafe(f)) { + Ok(Ok(value)) => Ok(value), + Ok(Err(error)) => { + set_error_slot(out_error, error.message); + Err(error.status) + } + Err(payload) => { + set_error_slot(out_error, panic_message(payload)); + Err(tikv_status_t::TIKV_STATUS_PANIC) + } + } +} + +fn empty_owned_bytes() -> tikv_owned_bytes_t { + tikv_owned_bytes_t { + ptr: ptr::null_mut(), + len: 0, + } +} + +fn empty_kv_pairs() -> tikv_kv_pairs_t { + tikv_kv_pairs_t { + ptr: ptr::null_mut(), + len: 0, + } +} + +fn boxed_future(future: F) -> PinBoxRequestFuture +where + F: std::future::Future> + Send + 'static, +{ + Box::pin(future) +} + +fn lock_completion(request: &RequestState) -> std::sync::MutexGuard<'_, CompletionState> { + request + .completion + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +fn request_ptr(raw: *mut tikv_async_request_t) -> *mut RequestState { + raw as *mut RequestState +} + +fn request_ref<'a>(raw: *const tikv_async_request_t) -> FfiResult<&'a RequestState> { + if raw.is_null() { + return Err(FfiError::invalid_argument("request must not be null")); + } + Ok(unsafe { &*(raw as *const RequestState) }) +} + +fn client_ref<'a>(raw: *const tikv_async_client_t) -> FfiResult<&'a AsyncClientHandle> { + if raw.is_null() { + return Err(FfiError::invalid_argument("client must not be null")); + } + Ok(unsafe { &*(raw as *const AsyncClientHandle) }) +} + +fn retain_request(request: *mut RequestState) { + unsafe { + (*request).retain_count.fetch_add(1, Ordering::AcqRel); + } +} + +fn release_request(request: *mut RequestState) { + let previous = unsafe { (*request).retain_count.fetch_sub(1, Ordering::AcqRel) }; + if previous == 1 { + unsafe { + drop(Box::from_raw(request)); + } + } +} + +fn inflight_add(core: &AsyncClientCore) { + let mut drain = core + .lifecycle + .drain + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + drain.inflight += 1; +} + +fn inflight_done(core: &AsyncClientCore) { + let mut drain = core + .lifecycle + .drain + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + drain.inflight -= 1; + if drain.inflight == 0 { + core.lifecycle.drain_ready.notify_all(); + } +} + +fn wait_drain(core: &AsyncClientCore) { + let mut drain = core + .lifecycle + .drain + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + while drain.inflight != 0 { + drain = core + .lifecycle + .drain_ready + .wait(drain) + .unwrap_or_else(|poisoned| poisoned.into_inner()); + } +} + +fn bytes_to_owned(bytes: Vec) -> tikv_owned_bytes_t { + if bytes.is_empty() { + return empty_owned_bytes(); + } + + let mut boxed = bytes.into_boxed_slice(); + let owned = tikv_owned_bytes_t { + ptr: boxed.as_mut_ptr(), + len: boxed.len(), + }; + mem::forget(boxed); + owned +} + +fn pairs_to_owned(pairs: Vec) -> tikv_kv_pairs_t { + if pairs.is_empty() { + return empty_kv_pairs(); + } + + let ffi_pairs: Vec = pairs + .into_iter() + .map(|pair| { + let (key, value): (Key, Vec) = pair.into(); + tikv_kv_pair_t { + key: bytes_to_owned(key.into()), + value: bytes_to_owned(value), + } + }) + .collect(); + let mut ffi_pairs = ffi_pairs.into_boxed_slice(); + let owned = tikv_kv_pairs_t { + ptr: ffi_pairs.as_mut_ptr(), + len: ffi_pairs.len(), + }; + mem::forget(ffi_pairs); + owned +} + +fn slice_is_empty(raw: tikv_slice_t) -> bool { + raw.len == 0 +} + +fn slice_to_vec(raw: tikv_slice_t, name: &str) -> FfiResult> { + if raw.len == 0 { + return Ok(Vec::new()); + } + if raw.ptr.is_null() { + return Err(FfiError::invalid_argument(format!( + "{name} pointer is null while length is {}", + raw.len + ))); + } + + let bytes = unsafe { slice::from_raw_parts(raw.ptr, raw.len) }; + Ok(bytes.to_vec()) +} + +fn slices_to_vec<'a, T>(ptr: *const T, len: usize, name: &str) -> FfiResult<&'a [T]> { + if len == 0 { + return Ok(&[]); + } + if ptr.is_null() { + return Err(FfiError::invalid_argument(format!( + "{name} pointer is null while length is {len}" + ))); + } + + Ok(unsafe { slice::from_raw_parts(ptr, len) }) +} + +fn slice_to_key(raw: tikv_slice_t, name: &str) -> FfiResult { + Ok(Key::from(slice_to_vec(raw, name)?)) +} + +fn slices_to_keys(raw: tikv_slices_t, name: &str) -> FfiResult> { + let slices = slices_to_vec(raw.ptr, raw.len, name)?; + slices + .iter() + .enumerate() + .map(|(index, item)| slice_to_key(*item, &format!("{name}[{index}]"))) + .collect() +} + +fn slice_mutations_to_owned(raw: tikv_async_raw_mutations_t) -> FfiResult> { + let mutations = slices_to_vec(raw.ptr, raw.len, "mutations")?; + mutations + .iter() + .enumerate() + .map(|(index, mutation)| { + let op = MutationOp::from_ffi(mutation.op, &format!("mutations[{index}]"))?; + let key = slice_to_key(mutation.key, &format!("mutations[{index}].key"))?; + match op { + MutationOp::Put => { + let value = slice_to_vec(mutation.value, &format!("mutations[{index}].value"))?; + Ok(OwnedMutation::Put(KvPair::new(key, value), 0)) + } + MutationOp::Delete => Ok(OwnedMutation::Delete(key)), + } + }) + .collect() +} + +fn slice_mutations_with_ttl_to_owned( + raw: tikv_async_raw_mutations_with_ttl_t, +) -> FfiResult> { + let mutations = slices_to_vec(raw.ptr, raw.len, "mutations")?; + mutations + .iter() + .enumerate() + .map(|(index, mutation)| { + let op = MutationOp::from_ffi(mutation.op, &format!("mutations[{index}]"))?; + let key = slice_to_key(mutation.key, &format!("mutations[{index}].key"))?; + match op { + MutationOp::Put => { + let value = slice_to_vec(mutation.value, &format!("mutations[{index}].value"))?; + Ok(OwnedMutation::Put( + KvPair::new(key, value), + mutation.ttl_secs, + )) + } + MutationOp::Delete => Ok(OwnedMutation::Delete(key)), + } + }) + .collect() +} + +fn build_write_batch_runs(mutations: Vec) -> Vec { + let mut runs = Vec::new(); + let mut put_pairs = Vec::new(); + let mut put_keys = HashSet::new(); + let mut delete_keys = Vec::new(); + + for mutation in mutations { + match mutation { + OwnedMutation::Put(pair, ttl_secs) => { + if !delete_keys.is_empty() { + runs.push(WriteBatchRun::Delete(mem::take(&mut delete_keys))); + } + + let key = pair.key().clone(); + if put_keys.contains(&key) && !put_pairs.is_empty() { + runs.push(WriteBatchRun::Put(mem::take(&mut put_pairs))); + put_keys.clear(); + } + put_keys.insert(key); + put_pairs.push((pair, ttl_secs)); + } + OwnedMutation::Delete(key) => { + if !put_pairs.is_empty() { + runs.push(WriteBatchRun::Put(mem::take(&mut put_pairs))); + put_keys.clear(); + } + delete_keys.push(key); + } + } + } + + if !put_pairs.is_empty() { + runs.push(WriteBatchRun::Put(put_pairs)); + } + if !delete_keys.is_empty() { + runs.push(WriteBatchRun::Delete(delete_keys)); + } + + runs +} + +fn slice_to_range(start_key: tikv_slice_t, end_key: tikv_slice_t) -> FfiResult { + let start = slice_to_key(start_key, "start_key")?; + let end = if slice_is_empty(end_key) { + None + } else { + Some(slice_to_key(end_key, "end_key")?) + }; + Ok(BoundRange::from((start, end))) +} + +fn parse_pd_endpoints(pd_endpoints_csv: *const c_char) -> FfiResult> { + if pd_endpoints_csv.is_null() { + return Err(FfiError::invalid_argument( + "pd_endpoints_csv must not be null", + )); + } + + let csv = unsafe { CStr::from_ptr(pd_endpoints_csv) }; + let csv = csv + .to_str() + .map_err(|_| FfiError::invalid_argument("pd_endpoints_csv must be valid UTF-8"))?; + + let endpoints: Vec = csv + .split(',') + .map(str::trim) + .filter(|item| !item.is_empty()) + .map(ToOwned::to_owned) + .collect(); + + if endpoints.is_empty() { + return Err(FfiError::invalid_argument( + "pd_endpoints_csv must contain at least one endpoint", + )); + } + + Ok(endpoints) +} + +fn build_config( + timeout_secs: u32, + use_default_keyspace: bool, + grpc_max_decoding_message_size: usize, +) -> Config { + let mut config = Config::default(); + if timeout_secs > 0 { + config = config.with_timeout(Duration::from_secs(timeout_secs.into())); + } + if use_default_keyspace { + config = config.with_default_keyspace(); + } + if grpc_max_decoding_message_size > 0 { + config = config.with_grpc_max_decoding_message_size(grpc_max_decoding_message_size); + } + config +} + +fn make_runtime() -> FfiResult { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|error| FfiError::error(error.to_string())) +} + +fn new_client_handle( + pd_endpoints_csv: *const c_char, + timeout_secs: u32, + use_default_keyspace: bool, + grpc_max_decoding_message_size: usize, +) -> FfiResult { + let endpoints = parse_pd_endpoints(pd_endpoints_csv)?; + let config = build_config( + timeout_secs, + use_default_keyspace, + grpc_max_decoding_message_size, + ); + let runtime = make_runtime()?; + let client = runtime + .block_on(RawClient::new_with_config(endpoints, config)) + .map_err(|error| FfiError::error(error.to_string()))?; + + Ok(AsyncClientHandle { + core: Arc::new(AsyncClientCore { + runtime, + client, + lifecycle: ClientLifecycle { + accepting: AtomicBool::new(true), + cancel_inflight: AtomicBool::new(false), + drain: Mutex::new(DrainState { inflight: 0 }), + drain_ready: Condvar::new(), + requests: Mutex::new(Vec::new()), + }, + }), + }) +} + +fn ensure_accepting(core: &AsyncClientCore) -> FfiResult<()> { + if !core.lifecycle.accepting.load(Ordering::Acquire) { + return Err(FfiError::rejected( + "client is shutting down and rejects new requests", + )); + } + Ok(()) +} + +fn ensure_request_kind(request: &RequestState, expected: RequestKind) -> FfiResult<()> { + if request.kind != expected { + return Err(FfiError::invalid_argument(format!( + "request does not carry a {} result", + expected.name() + ))); + } + Ok(()) +} + +fn make_request_state( + kind: RequestKind, + core: &Arc, + callback: tikv_async_completion_cb, + userdata: *mut c_void, +) -> *mut RequestState { + inflight_add(core); + let request = Box::new(RequestState { + kind, + completion: Mutex::new(CompletionState { + complete: false, + callback_running: false, + callback_done: false, + completion_status: None, + error_message: None, + output: None, + output_taken: false, + cancel_requested: false, + }), + ready: Condvar::new(), + client: Arc::downgrade(core), + callback: Some(callback), + userdata, + retain_count: AtomicUsize::new(1), + }); + let request_ptr = Box::into_raw(request); + let mut requests = core + .lifecycle + .requests + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + requests.push(request_ptr as usize); + request_ptr +} + +fn remove_request_from_client(core: &AsyncClientCore, request_ptr: *mut RequestState) { + let mut requests = core + .lifecycle + .requests + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + requests.retain(|entry| *entry != request_ptr as usize); +} + +fn finalize_request( + request_ptr: *mut RequestState, + status: tikv_status_t, + output: Option, + error_message: Option, +) { + let (callback, userdata) = { + let request = unsafe { &*request_ptr }; + let mut completion = lock_completion(request); + completion.complete = true; + completion.callback_running = true; + completion.completion_status = Some(status); + completion.output = output; + completion.error_message = + error_message.and_then(|message| CString::new(message.replace('\0', " ")).ok()); + request.ready.notify_all(); + (request.callback, request.userdata) + }; + + retain_request(request_ptr); + if let Some(callback) = callback { + callback(request_ptr as *mut tikv_async_request_t, status, userdata); + } + + let request = unsafe { &*request_ptr }; + { + let mut completion = lock_completion(request); + completion.callback_running = false; + completion.callback_done = true; + request.ready.notify_all(); + } + + if let Some(core) = request.client.upgrade() { + remove_request_from_client(&core, request_ptr); + inflight_done(&core); + } + + release_request(request_ptr); + release_request(request_ptr); +} + +async fn submit_put_batch(client: &RawClient, puts: Vec<(KvPair, u64)>) -> FfiResult<()> { + if puts.is_empty() { + return Ok(()); + } + + let (pairs, ttls): (Vec<_>, Vec<_>) = puts.into_iter().unzip(); + client + .batch_put_with_ttl(pairs, ttls) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(()) +} + +async fn submit_delete_batch(client: &RawClient, keys: Vec) -> FfiResult<()> { + if keys.is_empty() { + return Ok(()); + } + + client + .batch_delete(keys) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(()) +} + +async fn execute_write_batch(client: RawClient, runs: Vec) -> FfiResult<()> { + for run in runs { + match run { + WriteBatchRun::Put(pairs) => submit_put_batch(&client, pairs).await?, + WriteBatchRun::Delete(keys) => submit_delete_batch(&client, keys).await?, + } + } + Ok(()) +} + +fn spawn_request(core: Arc, request_ptr: *mut RequestState, future: F) +where + F: std::future::Future> + Send + 'static, +{ + let handle: Handle = core.runtime.handle().clone(); + let keepalive = Arc::clone(&core); + let request_addr = request_ptr as usize; + + handle.spawn(async move { + let result = AssertUnwindSafe(future).catch_unwind().await; + let request_ptr = request_addr as *mut RequestState; + let cancelled = { + let request = unsafe { &*request_ptr }; + let completion = lock_completion(request); + completion.cancel_requested + || keepalive.lifecycle.cancel_inflight.load(Ordering::Acquire) + }; + + let (status, output, error_message) = match result { + Ok(Ok(_)) if cancelled => ( + tikv_status_t::TIKV_STATUS_CANCELLED, + None, + Some("request cancelled".to_owned()), + ), + Ok(Ok(output)) => (tikv_status_t::TIKV_STATUS_OK, Some(output), None), + Ok(Err(error)) => (error.status, None, Some(error.message)), + Err(payload) => ( + tikv_status_t::TIKV_STATUS_PANIC, + None, + Some(panic_message(payload)), + ), + }; + + finalize_request(request_ptr, status, output, error_message); + drop(keepalive); + }); +} + +fn submit_common( + client: *const tikv_async_client_t, + kind: RequestKind, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + build_future: F, +) -> FfiResult<()> +where + F: FnOnce(RawClient) -> FfiResult, +{ + if out_request.is_null() { + return Err(FfiError::invalid_argument("out_request must not be null")); + } + + let completion = completion + .ok_or_else(|| FfiError::invalid_argument("completion callback must not be null"))?; + let handle = client_ref(client)?; + ensure_accepting(&handle.core)?; + + let request_ptr = make_request_state(kind, &handle.core, completion, userdata); + let future = match build_future(handle.core.client.clone()) { + Ok(future) => future, + Err(error) => { + remove_request_from_client(&handle.core, request_ptr); + inflight_done(&handle.core); + release_request(request_ptr); + return Err(error); + } + }; + + if !handle.core.lifecycle.accepting.load(Ordering::Acquire) { + remove_request_from_client(&handle.core, request_ptr); + inflight_done(&handle.core); + release_request(request_ptr); + return Err(FfiError::rejected( + "client is shutting down and rejects new requests", + )); + } + + retain_request(request_ptr); + unsafe { + *out_request = request_ptr as *mut tikv_async_request_t; + } + spawn_request(Arc::clone(&handle.core), request_ptr, future); + + if handle + .core + .lifecycle + .cancel_inflight + .load(Ordering::Acquire) + { + let request = unsafe { &*request_ptr }; + let mut completion = lock_completion(request); + completion.cancel_requested = true; + } + + Ok(()) +} + +fn take_request_output<'a>( + request: &'a RequestState, + expected: RequestKind, +) -> FfiResult> { + ensure_request_kind(request, expected)?; + let completion = lock_completion(request); + if !completion.complete { + return Err(FfiError::invalid_argument("request is still pending")); + } + if completion.completion_status != Some(tikv_status_t::TIKV_STATUS_OK) { + return Err(FfiError::invalid_argument( + "request did not complete successfully", + )); + } + if completion.output_taken { + return Err(FfiError::invalid_argument("request result already taken")); + } + Ok(completion) +} + +#[no_mangle] +pub extern "C" fn tikv_async_client_new( + pd_endpoints_csv: *const c_char, + timeout_secs: u32, + use_default_keyspace: bool, + out_client: *mut *mut tikv_async_client_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + tikv_async_client_new_with_grpc_max_decoding_message_size( + pd_endpoints_csv, + timeout_secs, + use_default_keyspace, + usize::MAX, + out_client, + out_error, + ) +} + +#[no_mangle] +pub extern "C" fn tikv_async_client_new_with_grpc_max_decoding_message_size( + pd_endpoints_csv: *const c_char, + timeout_secs: u32, + use_default_keyspace: bool, + grpc_max_decoding_message_size: usize, + out_client: *mut *mut tikv_async_client_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_client.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_client must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_client = ptr::null_mut(); + } + + match run_ffi(out_error, || { + let handle = new_client_handle( + pd_endpoints_csv, + timeout_secs, + use_default_keyspace, + grpc_max_decoding_message_size, + )?; + unsafe { + *out_client = Box::into_raw(Box::new(handle)) as *mut tikv_async_client_t; + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_client_shutdown_begin( + client: *mut tikv_async_client_t, + cancel_inflight: bool, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + match run_ffi(out_error, || { + let handle = client_ref(client)?; + handle + .core + .lifecycle + .accepting + .store(false, Ordering::Release); + if cancel_inflight { + handle + .core + .lifecycle + .cancel_inflight + .store(true, Ordering::Release); + let requests = handle + .core + .lifecycle + .requests + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(); + for request_addr in requests { + let request = unsafe { &*(request_addr as *const RequestState) }; + let mut completion = lock_completion(request); + completion.cancel_requested = true; + } + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_client_drain( + client: *const tikv_async_client_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + match run_ffi(out_error, || { + let handle = client_ref(client)?; + wait_drain(&handle.core); + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_put_submit( + client: *const tikv_async_client_t, + key: tikv_slice_t, + value: tikv_slice_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let key = slice_to_key(key, "key")?; + let value = slice_to_vec(value, "value")?; + submit_common( + client, + RequestKind::Put, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + client + .put(key, value) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Unit) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_put_with_ttl_submit( + client: *const tikv_async_client_t, + key: tikv_slice_t, + value: tikv_slice_t, + ttl_secs: u64, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let key = slice_to_key(key, "key")?; + let value = slice_to_vec(value, "value")?; + submit_common( + client, + RequestKind::PutWithTtl, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + client + .put_with_ttl(key, value, ttl_secs) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Unit) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_get_submit( + client: *const tikv_async_client_t, + key: tikv_slice_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let key = slice_to_key(key, "key")?; + submit_common( + client, + RequestKind::Get, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + let value = client + .get(key) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Get(value)) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_get_key_ttl_submit( + client: *const tikv_async_client_t, + key: tikv_slice_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let key = slice_to_key(key, "key")?; + submit_common( + client, + RequestKind::GetKeyTtl, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + let ttl = client + .get_key_ttl_secs(key) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::KeyTtl(ttl)) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_delete_submit( + client: *const tikv_async_client_t, + key: tikv_slice_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let key = slice_to_key(key, "key")?; + submit_common( + client, + RequestKind::Delete, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + client + .delete(key) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Unit) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_delete_range_submit( + client: *const tikv_async_client_t, + start_key: tikv_slice_t, + end_key: tikv_slice_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let range = slice_to_range(start_key, end_key)?; + submit_common( + client, + RequestKind::DeleteRange, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + client + .delete_range(range) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Unit) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_scan_submit( + client: *const tikv_async_client_t, + start_key: tikv_slice_t, + end_key: tikv_slice_t, + limit: u32, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let range = slice_to_range(start_key, end_key)?; + submit_common( + client, + RequestKind::Scan, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + let pairs = client + .scan(range, limit) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Pairs(pairs)) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_scan_reverse_submit( + client: *const tikv_async_client_t, + start_key: tikv_slice_t, + end_key: tikv_slice_t, + limit: u32, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + if slice_is_empty(start_key) { + return Err(FfiError::invalid_argument( + "scan_reverse does not support empty start_key in client-rust", + )); + } + let range = slice_to_range(start_key, end_key)?; + submit_common( + client, + RequestKind::ScanReverse, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + let pairs = client + .scan_reverse(range, limit) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Pairs(pairs)) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_batch_write_submit( + client: *const tikv_async_client_t, + mutations: tikv_async_raw_mutations_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let runs = build_write_batch_runs(slice_mutations_to_owned(mutations)?); + submit_common( + client, + RequestKind::BatchWrite, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + execute_write_batch(client, runs).await?; + Ok(RequestOutput::Unit) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_batch_write_with_ttl_submit( + client: *const tikv_async_client_t, + mutations: tikv_async_raw_mutations_with_ttl_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let runs = build_write_batch_runs(slice_mutations_with_ttl_to_owned(mutations)?); + submit_common( + client, + RequestKind::BatchWrite, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + execute_write_batch(client, runs).await?; + Ok(RequestOutput::Unit) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_raw_batch_get_submit( + client: *const tikv_async_client_t, + keys: tikv_slices_t, + completion: Option, + userdata: *mut c_void, + out_request: *mut *mut tikv_async_request_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_request.is_null() { + unsafe { *out_request = ptr::null_mut() }; + } + + match run_ffi(out_error, || { + let keys = slices_to_keys(keys, "keys")?; + submit_common( + client, + RequestKind::BatchGet, + completion, + userdata, + out_request, + move |client| { + Ok(boxed_future(async move { + if keys.is_empty() { + return Ok(RequestOutput::Pairs(Vec::new())); + } + let pairs = client + .batch_get(keys) + .await + .map_err(|error| FfiError::error(error.to_string()))?; + Ok(RequestOutput::Pairs(pairs)) + })) + }, + ) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_cancel( + request: *mut tikv_async_request_t, + out_cancel_requested: *mut bool, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_cancel_requested.is_null() { + clear_error_slot(out_error); + set_error_slot( + out_error, + "out_cancel_requested must not be null".to_owned(), + ); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_cancel_requested = false; + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let mut completion = lock_completion(request); + if completion.complete || completion.cancel_requested { + return Ok(()); + } + completion.cancel_requested = true; + unsafe { + *out_cancel_requested = true; + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_is_complete( + request: *const tikv_async_request_t, + out_complete: *mut bool, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_complete.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_complete must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_complete = false; + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let completion = lock_completion(request); + unsafe { + *out_complete = completion.complete; + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_completion_status( + request: *const tikv_async_request_t, + out_status: *mut tikv_status_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_status.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_status must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_status = tikv_status_t::TIKV_STATUS_ERROR; + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let completion = lock_completion(request); + if !completion.complete { + return Err(FfiError::invalid_argument("request is still pending")); + } + unsafe { + *out_status = completion.completion_status.expect("completed request"); + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_error_message( + request: *const tikv_async_request_t, +) -> *const c_char { + let Ok(request) = request_ref(request) else { + return ptr::null(); + }; + let completion = lock_completion(request); + completion + .error_message + .as_ref() + .map(|message| message.as_ptr()) + .unwrap_or(ptr::null()) +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_take_get_result( + request: *const tikv_async_request_t, + out_value: *mut tikv_owned_bytes_t, + out_found: *mut bool, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_value.is_null() || out_found.is_null() { + clear_error_slot(out_error); + set_error_slot( + out_error, + "out_value and out_found must not be null".to_owned(), + ); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_value = empty_owned_bytes(); + *out_found = false; + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let mut completion = take_request_output(request, RequestKind::Get)?; + match completion.output.as_mut().expect("successful request") { + RequestOutput::Get(value) => { + let value = value.take(); + unsafe { + match value { + Some(value) => { + *out_value = bytes_to_owned(value); + *out_found = true; + } + None => { + *out_value = empty_owned_bytes(); + *out_found = false; + } + } + } + completion.output_taken = true; + Ok(()) + } + _ => Err(FfiError::error("internal get result mismatch")), + } + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_take_get_key_ttl_result( + request: *const tikv_async_request_t, + out_ttl_secs: *mut u64, + out_found: *mut bool, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_ttl_secs.is_null() || out_found.is_null() { + clear_error_slot(out_error); + set_error_slot( + out_error, + "out_ttl_secs and out_found must not be null".to_owned(), + ); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_ttl_secs = 0; + *out_found = false; + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let mut completion = take_request_output(request, RequestKind::GetKeyTtl)?; + match completion.output.as_mut().expect("successful request") { + RequestOutput::KeyTtl(ttl) => { + let ttl = ttl.take(); + unsafe { + match ttl { + Some(ttl) => { + *out_ttl_secs = ttl; + *out_found = true; + } + None => { + *out_ttl_secs = 0; + *out_found = false; + } + } + } + completion.output_taken = true; + Ok(()) + } + _ => Err(FfiError::error("internal get_key_ttl result mismatch")), + } + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_take_scan_result( + request: *const tikv_async_request_t, + out_pairs: *mut tikv_kv_pairs_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_pairs.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_pairs must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_pairs = empty_kv_pairs(); + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let expected = match request.kind { + RequestKind::Scan => RequestKind::Scan, + RequestKind::ScanReverse => RequestKind::ScanReverse, + _ => { + return Err(FfiError::invalid_argument( + "request does not carry a scan result", + )) + } + }; + let mut completion = take_request_output(request, expected)?; + match completion.output.take().expect("successful request") { + RequestOutput::Pairs(pairs) => { + unsafe { + *out_pairs = pairs_to_owned(pairs); + } + completion.output_taken = true; + Ok(()) + } + _ => Err(FfiError::error("internal scan result mismatch")), + } + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_take_batch_get_result( + request: *const tikv_async_request_t, + out_pairs: *mut tikv_kv_pairs_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_pairs.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_pairs must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_pairs = empty_kv_pairs(); + } + + match run_ffi(out_error, || { + let request = request_ref(request)?; + let mut completion = take_request_output(request, RequestKind::BatchGet)?; + match completion.output.take().expect("successful request") { + RequestOutput::Pairs(pairs) => { + unsafe { + *out_pairs = pairs_to_owned(pairs); + } + completion.output_taken = true; + Ok(()) + } + _ => Err(FfiError::error("internal batch_get result mismatch")), + } + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_metrics_snapshot_collect( + out_snapshot: *mut *mut tikv_metrics_snapshot_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_snapshot.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_snapshot must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + unsafe { + *out_snapshot = ptr::null_mut(); + } + + match run_ffi(out_error, || { + let snapshot = Box::new(collect_metrics_snapshot()); + unsafe { + *out_snapshot = Box::into_raw(snapshot) as *mut tikv_metrics_snapshot_t; + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_metrics_snapshot_len(snapshot: *const tikv_metrics_snapshot_t) -> usize { + if snapshot.is_null() { + return 0; + } + unsafe { (*(snapshot as *const MetricsSnapshot)).samples.len() } +} + +#[no_mangle] +pub extern "C" fn tikv_metrics_snapshot_get( + snapshot: *const tikv_metrics_snapshot_t, + index: usize, + out_sample: *mut tikv_metric_sample_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if out_sample.is_null() { + clear_error_slot(out_error); + set_error_slot(out_error, "out_sample must not be null".to_owned()); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + match run_ffi(out_error, || { + let snapshot = metrics_snapshot_ref(snapshot)?; + let sample = snapshot.samples.get(index).ok_or_else(|| { + FfiError::invalid_argument(format!("sample index {index} out of range")) + })?; + unsafe { + *out_sample = sample.ffi; + } + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[no_mangle] +pub extern "C" fn tikv_error_message(error: *const tikv_error_t) -> *const c_char { + if error.is_null() { + return ptr::null(); + } + unsafe { (*error).message.as_ptr() } +} + +#[no_mangle] +pub extern "C" fn tikv_error_destroy(error: *mut tikv_error_t) { + if error.is_null() { + return; + } + unsafe { + drop(Box::from_raw(error)); + } +} + +#[no_mangle] +pub extern "C" fn tikv_owned_bytes_destroy(bytes: tikv_owned_bytes_t) { + if bytes.ptr.is_null() || bytes.len == 0 { + return; + } + unsafe { + drop(Vec::from_raw_parts(bytes.ptr, bytes.len, bytes.len)); + } +} + +#[no_mangle] +pub extern "C" fn tikv_kv_pairs_destroy(pairs: tikv_kv_pairs_t) { + if pairs.ptr.is_null() || pairs.len == 0 { + return; + } + unsafe { + let pairs = Vec::from_raw_parts(pairs.ptr, pairs.len, pairs.len); + for pair in pairs { + tikv_owned_bytes_destroy(pair.key); + tikv_owned_bytes_destroy(pair.value); + } + } +} + +#[no_mangle] +pub extern "C" fn tikv_metrics_snapshot_destroy(snapshot: *mut tikv_metrics_snapshot_t) { + if snapshot.is_null() { + return; + } + unsafe { + drop(Box::from_raw(snapshot as *mut MetricsSnapshot)); + } +} + +#[no_mangle] +pub extern "C" fn tikv_async_request_destroy(request: *mut tikv_async_request_t) { + if request.is_null() { + return; + } + release_request(request_ptr(request)); +} + +#[no_mangle] +pub extern "C" fn tikv_async_client_destroy(client: *mut tikv_async_client_t) { + if client.is_null() { + return; + } + let handle = unsafe { Box::from_raw(client as *mut AsyncClientHandle) }; + handle + .core + .lifecycle + .accepting + .store(false, Ordering::Release); + handle + .core + .lifecycle + .cancel_inflight + .store(true, Ordering::Release); + let requests = handle + .core + .lifecycle + .requests + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(); + for request_addr in requests { + let request = unsafe { &*(request_addr as *const RequestState) }; + let mut completion = lock_completion(request); + completion.cancel_requested = true; + } + wait_drain(&handle.core); +} + +#[cfg(test)] +mod tests { + use super::*; + + extern "C" fn noop_completion( + _request: *mut tikv_async_request_t, + _status: tikv_status_t, + _userdata: *mut c_void, + ) { + } + + fn synthetic_request(kind: RequestKind, output: RequestOutput) -> *mut tikv_async_request_t { + Box::into_raw(Box::new(RequestState { + kind, + completion: Mutex::new(CompletionState { + complete: true, + callback_running: false, + callback_done: true, + completion_status: Some(tikv_status_t::TIKV_STATUS_OK), + error_message: None, + output: Some(output), + output_taken: false, + cancel_requested: false, + }), + ready: Condvar::new(), + client: Weak::new(), + callback: Some(noop_completion), + userdata: ptr::null_mut(), + retain_count: AtomicUsize::new(1), + })) as *mut tikv_async_request_t + } + + #[test] + fn parse_pd_endpoints_csv() { + let input = CString::new(" 127.0.0.1:2379 , 127.0.0.1:2380 ").unwrap(); + let endpoints = parse_pd_endpoints(input.as_ptr()).unwrap(); + assert_eq!(endpoints, vec!["127.0.0.1:2379", "127.0.0.1:2380"]); + } + + #[test] + fn ffi_client_grpc_decoding_limit_is_configurable() { + let default_limit = Config::default().grpc_max_decoding_message_size; + assert_eq!(default_limit, 4 * 1024 * 1024); + assert_eq!( + build_config(0, false, 0).grpc_max_decoding_message_size, + default_limit + ); + assert_eq!( + build_config(0, false, 32 * 1024 * 1024).grpc_max_decoding_message_size, + 32 * 1024 * 1024 + ); + assert_eq!( + build_config(0, false, usize::MAX).grpc_max_decoding_message_size, + usize::MAX + ); + } + + #[test] + fn metrics_snapshot_collects_supported_counter_and_histogram() { + let counter = prometheus::IntCounterVec::new( + prometheus::Opts::new("tikv_request_total", "test tikv request total"), + &["type"], + ) + .expect("counter vec"); + prometheus::default_registry() + .register(Box::new(counter.clone())) + .expect("register counter"); + counter.with_label_values(&["raw_get"]).inc_by(3); + + let histogram = prometheus::HistogramVec::new( + prometheus::HistogramOpts::new( + "tikv_request_duration_seconds", + "test tikv request duration", + ), + &["type"], + ) + .expect("histogram vec"); + prometheus::default_registry() + .register(Box::new(histogram.clone())) + .expect("register histogram"); + histogram.with_label_values(&["raw_get"]).observe(0.01); + + let mut error = ptr::null_mut(); + let mut snapshot = ptr::null_mut(); + let status = tikv_metrics_snapshot_collect(&mut snapshot, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert!(!snapshot.is_null()); + + let len = tikv_metrics_snapshot_len(snapshot); + assert!(len > 0); + + let mut raw_get_total = None; + let mut raw_get_duration_count = None; + let mut saw_inf_bucket = false; + + for index in 0..len { + let mut sample = tikv_metric_sample_t { + name: ptr::null(), + type_label: ptr::null(), + kind: tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_UNKNOWN, + value: 0.0, + le: 0.0, + }; + let status = tikv_metrics_snapshot_get(snapshot, index, &mut sample, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + + let name = unsafe { CStr::from_ptr(sample.name) }.to_str().unwrap(); + let type_label = if sample.type_label.is_null() { + None + } else { + Some( + unsafe { CStr::from_ptr(sample.type_label) } + .to_str() + .unwrap(), + ) + }; + + if name == "tikv_request_total" + && sample.kind == tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_COUNTER + && type_label == Some("raw_get") + { + raw_get_total = Some(sample.value); + } + + if name == "tikv_request_duration_seconds" + && sample.kind == tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT + && type_label == Some("raw_get") + { + raw_get_duration_count = Some(sample.value); + } + + if name == "tikv_request_duration_seconds" + && sample.kind == tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET + && type_label == Some("raw_get") + && sample.le.is_infinite() + { + saw_inf_bucket = true; + } + } + + assert_eq!(raw_get_total, Some(3.0)); + assert_eq!(raw_get_duration_count, Some(1.0)); + assert!(saw_inf_bucket); + + tikv_metrics_snapshot_destroy(snapshot); + } + + #[test] + fn write_batch_runs_preserve_put_delete_phase_order() { + let runs = build_write_batch_runs(vec![ + OwnedMutation::Put(KvPair::new(b"a".to_vec(), b"v1".to_vec()), 10), + OwnedMutation::Put(KvPair::new(b"b".to_vec(), b"v2".to_vec()), 0), + OwnedMutation::Delete(Key::from(b"c".to_vec())), + OwnedMutation::Delete(Key::from(b"d".to_vec())), + OwnedMutation::Put(KvPair::new(b"e".to_vec(), b"v3".to_vec()), 30), + ]); + + assert_eq!( + runs, + vec![ + WriteBatchRun::Put(vec![ + (KvPair::new(b"a".to_vec(), b"v1".to_vec()), 10), + (KvPair::new(b"b".to_vec(), b"v2".to_vec()), 0), + ]), + WriteBatchRun::Delete(vec![Key::from(b"c".to_vec()), Key::from(b"d".to_vec())]), + WriteBatchRun::Put(vec![(KvPair::new(b"e".to_vec(), b"v3".to_vec()), 30,)]), + ] + ); + } + + #[test] + fn write_batch_splits_repeated_put_keys() { + let runs = build_write_batch_runs(vec![ + OwnedMutation::Put(KvPair::new(b"a".to_vec(), b"v1".to_vec()), 10), + OwnedMutation::Put(KvPair::new(b"b".to_vec(), b"v2".to_vec()), 20), + OwnedMutation::Put(KvPair::new(b"a".to_vec(), b"v3".to_vec()), 30), + ]); + + assert_eq!( + runs, + vec![ + WriteBatchRun::Put(vec![ + (KvPair::new(b"a".to_vec(), b"v1".to_vec()), 10), + (KvPair::new(b"b".to_vec(), b"v2".to_vec()), 20), + ]), + WriteBatchRun::Put(vec![(KvPair::new(b"a".to_vec(), b"v3".to_vec()), 30,)]), + ] + ); + } + + #[test] + fn ffi_mutations_copy_put_ttl_and_ignore_delete_ttl() { + let put_key = b"put-key"; + let put_value = b"put-value"; + let delete_key = b"delete-key"; + let mutations = [ + tikv_async_raw_mutation_with_ttl_t { + op: TIKV_ASYNC_RAW_MUTATION_OP_PUT, + key: tikv_slice_t { + ptr: put_key.as_ptr(), + len: put_key.len(), + }, + value: tikv_slice_t { + ptr: put_value.as_ptr(), + len: put_value.len(), + }, + ttl_secs: 42, + }, + tikv_async_raw_mutation_with_ttl_t { + op: TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + key: tikv_slice_t { + ptr: delete_key.as_ptr(), + len: delete_key.len(), + }, + value: tikv_slice_t { + ptr: ptr::null(), + len: 0, + }, + ttl_secs: u64::MAX, + }, + ]; + + let owned = slice_mutations_with_ttl_to_owned(tikv_async_raw_mutations_with_ttl_t { + ptr: mutations.as_ptr(), + len: mutations.len(), + }) + .unwrap(); + + assert_eq!( + owned, + vec![ + OwnedMutation::Put(KvPair::new(put_key.to_vec(), put_value.to_vec()), 42,), + OwnedMutation::Delete(Key::from(delete_key.to_vec())), + ] + ); + } + + #[test] + fn legacy_ffi_mutations_default_put_ttl_to_zero() { + let key = b"legacy-key"; + let value = b"legacy-value"; + let mutations = [tikv_async_raw_mutation_t { + op: TIKV_ASYNC_RAW_MUTATION_OP_PUT, + key: tikv_slice_t { + ptr: key.as_ptr(), + len: key.len(), + }, + value: tikv_slice_t { + ptr: value.as_ptr(), + len: value.len(), + }, + }]; + + let owned = slice_mutations_to_owned(tikv_async_raw_mutations_t { + ptr: mutations.as_ptr(), + len: mutations.len(), + }) + .unwrap(); + + assert_eq!( + owned, + vec![OwnedMutation::Put( + KvPair::new(key.to_vec(), value.to_vec()), + 0, + )] + ); + } + + #[test] + fn request_get_result_transfers_owned_bytes_once() { + let request = synthetic_request( + RequestKind::Get, + RequestOutput::Get(Some(b"value".to_vec())), + ); + + let mut error = ptr::null_mut(); + let mut out_value = empty_owned_bytes(); + let mut out_found = false; + let status = + tikv_async_request_take_get_result(request, &mut out_value, &mut out_found, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert!(out_found); + let bytes = unsafe { slice::from_raw_parts(out_value.ptr, out_value.len) }; + assert_eq!(bytes, b"value"); + tikv_owned_bytes_destroy(out_value); + + let status = + tikv_async_request_take_get_result(request, &mut out_value, &mut out_found, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(!error.is_null()); + tikv_error_destroy(error); + tikv_async_request_destroy(request); + } + + #[test] + fn request_get_key_ttl_result_distinguishes_found_states() { + for (ttl, expected_found, expected_ttl) in + [(Some(42), true, 42), (Some(0), true, 0), (None, false, 0)] + { + let request = synthetic_request(RequestKind::GetKeyTtl, RequestOutput::KeyTtl(ttl)); + + let mut error = ptr::null_mut(); + let mut out_ttl_secs = u64::MAX; + let mut out_found = false; + let status = tikv_async_request_take_get_key_ttl_result( + request, + &mut out_ttl_secs, + &mut out_found, + &mut error, + ); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert_eq!(out_found, expected_found); + assert_eq!(out_ttl_secs, expected_ttl); + + let status = tikv_async_request_take_get_key_ttl_result( + request, + &mut out_ttl_secs, + &mut out_found, + &mut error, + ); + assert_eq!(status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(!error.is_null()); + tikv_error_destroy(error); + tikv_async_request_destroy(request); + } + } + + #[test] + fn request_get_key_ttl_result_rejects_wrong_request_kind() { + let request = synthetic_request(RequestKind::Get, RequestOutput::Get(None)); + let mut error = ptr::null_mut(); + let mut out_ttl_secs = u64::MAX; + let mut out_found = true; + + let status = tikv_async_request_take_get_key_ttl_result( + request, + &mut out_ttl_secs, + &mut out_found, + &mut error, + ); + + assert_eq!(status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(!error.is_null()); + assert_eq!(out_ttl_secs, 0); + assert!(!out_found); + tikv_error_destroy(error); + tikv_async_request_destroy(request); + } + + #[test] + fn request_batch_get_result_transfers_pairs_once() { + let request = synthetic_request( + RequestKind::BatchGet, + RequestOutput::Pairs(vec![KvPair::new(b"k1".to_vec(), b"v1".to_vec())]), + ); + + let mut error = ptr::null_mut(); + let mut out_pairs = empty_kv_pairs(); + let status = tikv_async_request_take_batch_get_result(request, &mut out_pairs, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert_eq!(out_pairs.len, 1); + let key = + unsafe { slice::from_raw_parts((*out_pairs.ptr).key.ptr, (*out_pairs.ptr).key.len) }; + let value = unsafe { + slice::from_raw_parts((*out_pairs.ptr).value.ptr, (*out_pairs.ptr).value.len) + }; + assert_eq!(key, b"k1"); + assert_eq!(value, b"v1"); + tikv_kv_pairs_destroy(out_pairs); + + let status = tikv_async_request_take_batch_get_result(request, &mut out_pairs, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(!error.is_null()); + tikv_error_destroy(error); + tikv_async_request_destroy(request); + } + + #[test] + fn request_cancel_is_idempotent() { + let request = Box::into_raw(Box::new(RequestState { + kind: RequestKind::Put, + completion: Mutex::new(CompletionState { + complete: false, + callback_running: false, + callback_done: false, + completion_status: None, + error_message: None, + output: None, + output_taken: false, + cancel_requested: false, + }), + ready: Condvar::new(), + client: Weak::new(), + callback: Some(noop_completion), + userdata: ptr::null_mut(), + retain_count: AtomicUsize::new(1), + })) as *mut tikv_async_request_t; + + let mut error = ptr::null_mut(); + let mut cancel_requested = false; + let status = tikv_async_request_cancel(request, &mut cancel_requested, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert!(cancel_requested); + + let status = tikv_async_request_cancel(request, &mut cancel_requested, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert!(!cancel_requested); + + tikv_async_request_destroy(request); + } +}