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/README.md b/README.md index 619e683b..cb641781 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,8 @@ This crate provides an easy-to-use client for [TiKV](https://github.com/tikv/tik This crate lets you connect to a TiKV(>= `v5.0.0`) cluster and use either a transactional or raw (simple get/put style without transactional consistency guarantees) API to access and update your data. +For C applications, see the [asynchronous C binding](bindings/c/README.md). + The TiKV Rust client is an open source (Apache 2) project maintained by the TiKV Authors. We welcome contributions, see below for more info. Note that the current release is not suitable for production use - APIs are not yet stable and the crate has not been thoroughly tested in real-life use. 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..7cc5bb73 --- /dev/null +++ b/bindings/c/README.md @@ -0,0 +1,97 @@ +# TiKV Rust Client C binding + +This workspace crate exposes an asynchronous C ABI for TiKV RawKV operations. +The public headers are [`tikv_client_async_c.h`](include/tikv_client_async_c.h) +and [`tikv_client_metrics_c.h`](include/tikv_client_metrics_c.h). + +## Build + +From the repository root: + +```sh +cargo build -p tikv-client-rust-c +``` + +The build produces `target/debug/libtikv_client_rust_c.a` and a shared library +(`.so` on Linux or `.dylib` on macOS). Include headers from +`bindings/c/include` and link against `tikv_client_rust_c`. + +## API overview + +The C ABI supports raw put, get, delete, delete-range, forward and reverse +scan, mixed put/delete batches, and batch-get. Submitting a request copies its +input buffers before returning. +An accepted request returns a handle and eventually invokes its completion +callback on a Rust runtime worker thread. The callback may read the result or +destroy the request handle. Destroy each request exactly once, either in its +callback or after synchronizing with the callback to observe completion; do not +destroy a pending request. A callback can run before submit returns, so the +returned handle must not be used if the callback already destroyed it. +Error and result accessors may be used only in the completion callback or after +synchronizing with it to observe completion. The callback's `final_status` is +the terminal request status. Serialize all accessors and destruction for the +same request, including calls from its callback. A cross-thread handoff must +happen after the previous thread's last access and must synchronize the two +threads. Different requests may be accessed concurrently. +Each returned buffer, error, and snapshot has a matching destroy function in the +headers. + +`tikv_async_client_new` leaves the gRPC response decoding limit unlimited for +compatibility. Use +`tikv_async_client_new_with_grpc_max_decoding_message_size` to set a limit in +bytes. Passing zero to the latter uses the Rust client's default (currently +4 MiB). + +`tikv_async_client_shutdown_begin` rejects new submissions and lets accepted +requests finish normally. `tikv_async_client_destroy` begins shutdown, waits for +all accepted callbacks to return, and releases the client. Stop other calls using +the client handle before destroying it. Do not destroy a client from one of its +own callbacks. + +An empty reverse-scan start key is unsupported by the underlying client. +Mixed write batches preserve the order of adjacent put and delete phases, but +may issue multiple TiKV requests. + +## Metrics snapshots + +`tikv_metrics_snapshot_collect` exports only the eight client-rust TiKV/PD +request counter and duration histogram families named in +`tikv_client_metrics_c.h`. It reads the process-wide Prometheus default registry, +so values include all client instances, including callers outside the C API. +It accepts only the expected metric type for each family and samples whose sole +label is `type`; other families and incompatible samples are skipped. Histogram +samples include cumulative buckets (including +Inf), sum, and count. One call +returns the snapshot handle and a read-only view of all samples in a contiguous +array. The array and its strings remain valid until the snapshot is destroyed +with `tikv_metrics_snapshot_destroy`. For an empty snapshot, the array view is +`{NULL, 0}`. + +The existing `tikv_request_total` and `pd_request_total` metrics count +underlying RPC attempts. A single C operation can issue multiple RPCs due to +retries, scans, or mixed write batches. + +## C example + +Build the library first, then compile the example from the repository root. + +macOS: + +```sh +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: + +```sh +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 +``` + +With a TiKV cluster and PD endpoint running: + +```sh +./target/debug/raw_async_demo 127.0.0.1:2379 +``` diff --git a/bindings/c/examples/raw_async_demo.c b/bindings/c/examples/raw_async_demo.c new file mode 100644 index 00000000..d65694a2 --- /dev/null +++ b/bindings/c/examples/raw_async_demo.c @@ -0,0 +1,314 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +#include "tikv_client_async_c.h" + +#include +#include +#include +#include +#include + +typedef struct request_latch_t { + atomic_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; + atomic_store_explicit(&latch->done, 1, memory_order_release); + (void)request; +} + +static int wait_latch_ok(const char *label, + tikv_async_request_t *request, + request_latch_t *latch) { + while (!atomic_load_explicit(&latch->done, memory_order_acquire)) { + 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_t write_items_buf[] = { + {.op = TIKV_ASYNC_RAW_MUTATION_OP_PUT, + .key = key_a, + .value = value_a_v1}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_PUT, + .key = key_b, + .value = value_b}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + .key = key_a, + .value = {NULL, 0}}, + {.op = TIKV_ASYNC_RAW_MUTATION_OP_PUT, + .key = key_a, + .value = value_a_v2}, + }; + const tikv_async_raw_mutations_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_client_destroy(client); + tikv_async_request_destroy(reset); + return 1; + } + tikv_async_request_destroy(reset); + + status = tikv_async_raw_batch_write_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_client_destroy(client); + tikv_async_request_destroy(batch_write); + 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_client_destroy(client); + tikv_async_request_destroy(batch_get); + 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_client_destroy(client); + tikv_async_request_destroy(batch_get); + tikv_async_request_destroy(scan); + 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_client_destroy(client); + tikv_async_request_destroy(batch_get); + tikv_async_request_destroy(scan); + tikv_async_request_destroy(reverse_scan); + 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_client_destroy(client); + tikv_async_request_destroy(batch_get); + tikv_async_request_destroy(scan); + tikv_async_request_destroy(reverse_scan); + 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_client_destroy(client); + tikv_async_request_destroy(scan); + tikv_async_request_destroy(reverse_scan); + 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_client_destroy(client); + tikv_async_request_destroy(reverse_scan); + 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_client_destroy(client); + tikv_async_request_destroy(cleanup); + 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_client_destroy(client); + tikv_async_request_destroy(final_scan); + 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_client_destroy(client); + tikv_async_request_destroy(final_scan); + 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, &error); + if (!expect_ok("shutdown_begin", status, error)) { + tikv_async_client_destroy(client); + return 1; + } + tikv_async_client_destroy(client); + return 0; +} 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..2a1c2901 --- /dev/null +++ b/bindings/c/include/tikv_client_async_c.h @@ -0,0 +1,376 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +#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 and request ownership explicit. + * + * 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 during client shutdown. + * - 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`. + * - Destroy the request handle exactly once, either in its completion callback + * or after the caller has synchronized with the callback to observe completion. + * Do not destroy a pending request before its callback starts. + * - A callback may run before the submit function returns. If the callback + * destroys the handle, the caller must not use the returned handle. + * - Error and result accessors may be called only in the completion callback + * or after synchronizing with it to observe completion. The callback's + * `final_status` is the terminal request status. + * - The caller must serialize all status/error/result accessor and destroy + * calls for the same request, including calls from its callback. Before + * handing the request to another thread, finish accessing it and synchronize + * the handoff; the callback must not access it again after that handoff. + * Different request handles may be accessed concurrently. + * - 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 and + * lets accepted requests finish normally. + * - `tikv_async_client_destroy` begins shutdown, waits for all accepted requests + * to finish their callbacks, and then releases the client. + * - Do not call `tikv_async_client_destroy` from a completion callback belonging + * to the same client, or you may deadlock. + * - Before destroying a client, stop all other calls using its handle. + */ + +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; + +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_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_PANIC`: Rust caught a panic while executing the request. + * + * The callback may call request result accessors 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. + * + * Accepted requests finish normally and retain their operation status. + */ +tikv_status_t tikv_async_client_shutdown_begin(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 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 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 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. + */ + +/* + * Return the terminal request error message, or NULL if none. + * + * This may be called only after completion. + * 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); + +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); + +/* + * Release a completed request handle. NULL is a no-op. + * + * Call exactly once, in the completion callback or after synchronizing with it. + * Destroying a pending request before its callback starts is invalid. + */ +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_c.h b/bindings/c/include/tikv_client_metrics_c.h new file mode 100644 index 00000000..7eb60427 --- /dev/null +++ b/bindings/c/include/tikv_client_metrics_c.h @@ -0,0 +1,101 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +#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 families exported by `tikv_metrics_snapshot_collect`. + * Request totals are counters; request durations are histograms. + */ +#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" + +/* + * Values of the Prometheus label key `type` used by the C RawKV operations + * and their PD lookups. Other `type` values may appear when callers outside + * the C API use these same metric families. + */ +#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_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_DELETE_RANGE "raw_delete_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_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; + +typedef struct tikv_metric_samples_t { + const tikv_metric_sample_t *ptr; + size_t len; +} tikv_metric_samples_t; + +/* + * Collect a structured snapshot of only the eight client-rust TiKV/PD request + * metric families named above, from the process Prometheus default registry. + * Values are shared by all client instances in the process, including callers + * outside the C API; they are not scoped to a particular C client handle. + * + * Only counters for request totals and histograms for request durations are + * exported. Each sample must have exactly one label, named `type`. Other + * families, unexpected metric types, and samples with missing or additional + * labels are skipped. Histograms include cumulative buckets (including +Inf), + * sum, and count. + * + * On success, `out_samples` views every sample in one contiguous array owned + * by `out_snapshot`. The array and its strings remain valid until + * `tikv_metrics_snapshot_destroy`; callers must not modify or free them. + * For an empty snapshot, `*out_samples` is {NULL, 0}. + * `type_label` is always present and contains the Prometheus `type` label value. + * On failure, each non-NULL output pointer is cleared to NULL or {NULL, 0}. + */ +tikv_status_t tikv_metrics_snapshot_collect( + tikv_metrics_snapshot_t **out_snapshot, + tikv_metric_samples_t *out_samples, + 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/client.rs b/bindings/c/src/client.rs new file mode 100644 index 00000000..5327d8b6 --- /dev/null +++ b/bindings/c/src/client.rs @@ -0,0 +1,172 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +use super::*; + +pub(crate) 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) +} + +pub(crate) 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: Arc::new(ClientLifecycle { + accepting: AtomicBool::new(true), + inflight: Mutex::new(0), + drain_ready: Condvar::new(), + }), + }), + }) +} +#[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, + 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); + Ok(()) + }) { + Ok(()) => tikv_status_t::TIKV_STATUS_OK, + Err(status) => status, + } +} + +#[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); + wait_drain(&handle.core.lifecycle); +} diff --git a/bindings/c/src/lib.rs b/bindings/c/src/lib.rs new file mode 100644 index 00000000..9cc8e678 --- /dev/null +++ b/bindings/c/src/lib.rs @@ -0,0 +1,553 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +#![allow(non_camel_case_types)] + +mod client; +mod metrics; +mod raw; +mod request; +#[cfg(test)] +mod tests; + +pub use client::*; +pub use metrics::*; +pub use raw::*; +pub use request::*; + +use std::cell::{RefCell, RefMut}; +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, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +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, +} + +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, +} + +#[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, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct tikv_metric_samples_t { + pub ptr: *const tikv_metric_sample_t, + pub len: usize, +} + +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, + Get, + Delete, + DeleteRange, + Scan, + ScanReverse, + BatchWrite, + BatchGet, +} + +impl RequestKind { + fn name(self) -> &'static str { + match self { + Self::Put => "put", + Self::Get => "get", + 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), + Delete(Key), +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum WriteBatchRun { + Put(Vec), + Delete(Vec), +} + +#[derive(Debug)] +enum RequestOutput { + /// Successful completion with no return data, analogous to Rust's unit type `()`. + /// Used by put, delete, delete-range, and batch-write requests; the completion + /// callback reports their status, so no result extraction accessor is needed. + Unit, + Get(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 ClientLifecycle { + accepting: AtomicBool, + inflight: Mutex, + drain_ready: Condvar, +} + +/// Keeps a client operation counted until its callback has returned. +/// +/// This is separate from `RequestState` because C may retain a completed +/// request handle after the callback, while client shutdown only waits for +/// callbacks that are still running. +struct InflightGuard { + lifecycle: Arc, +} + +impl InflightGuard { + fn new(lifecycle: &Arc) -> Self { + inflight_add(lifecycle); + Self { + lifecycle: Arc::clone(lifecycle), + } + } +} + +impl Drop for InflightGuard { + /// Release the registration after the callback returns, or undo it when + /// shutdown rejects a submission. During shutdown, reaching zero wakes + /// `wait_drain` while sharing the counter mutex prevents a missed wakeup. + fn drop(&mut self) { + let mut inflight = self + .lifecycle + .inflight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *inflight -= 1; + if *inflight == 0 && !self.lifecycle.accepting.load(Ordering::Acquire) { + self.lifecycle.drain_ready.notify_all(); + } + } +} + +struct AsyncClientCore { + runtime: tokio::runtime::Runtime, + client: RawClient, + lifecycle: Arc, +} + +struct AsyncClientHandle { + core: Arc, +} + +struct CompletionState { + complete: bool, + completion_status: Option, + error_message: Option, + output: Option, + output_taken: bool, +} + +struct RequestState { + kind: RequestKind, + completion: RefCell, + callback: Option, + userdata: *mut c_void, +} + +// SAFETY: the C contract requires callback userdata to be usable on a runtime +// worker thread. +unsafe impl Send for RequestState {} +// SAFETY: the task finishes writing completion and releases its borrow before +// invoking the callback. The C contract requires synchronization with that +// callback before access on another thread, and serializes all accessors and +// destruction for the same request. No completion borrows overlap across threads. +unsafe impl Sync for RequestState {} + +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 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) }) +} + +/// Register work that client destruction must wait for before dropping its +/// runtime. Submitters and runtime workers update this per-client count +/// concurrently, so the mutex prevents lost increments/decrements. Increment +/// before spawning: even an immediately completed task must already be counted. +fn inflight_add(lifecycle: &ClientLifecycle) { + let mut inflight = lifecycle + .inflight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *inflight += 1; +} + +fn wait_drain(lifecycle: &ClientLifecycle) { + let mut inflight = lifecycle + .inflight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + while *inflight != 0 { + inflight = lifecycle + .drain_ready + .wait(inflight) + .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)?)) +} + +#[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); + } + } +} diff --git a/bindings/c/src/metrics.rs b/bindings/c/src/metrics.rs new file mode 100644 index 00000000..105419d1 --- /dev/null +++ b/bindings/c/src/metrics.rs @@ -0,0 +1,346 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +use super::*; + +struct MetricsSnapshot { + samples: Vec, + // The samples point into these CString buffers, which remain valid until + // the snapshot is destroyed even if this Vec reallocates. + _strings: Vec, +} + +fn ffi_cstring(value: &str) -> CString { + CString::new(value.replace('\0', " ")).expect("sanitized string") +} + +fn metric_type_label(metric: &prometheus::proto::Metric) -> Option<&str> { + match metric.get_label() { + [label] if label.get_name() == "type" => Some(label.get_value()), + _ => None, + } +} + +impl MetricsSnapshot { + fn push( + &mut self, + name: &str, + kind: tikv_metric_sample_kind_t, + type_label: &str, + value: f64, + le: f64, + ) { + let name = ffi_cstring(name); + let type_label = ffi_cstring(type_label); + self.samples.push(tikv_metric_sample_t { + name: name.as_ptr(), + type_label: type_label.as_ptr(), + kind, + value, + le, + }); + self._strings.push(name); + self._strings.push(type_label); + } +} + +fn collect_metrics_snapshot(registry: &prometheus::Registry) -> MetricsSnapshot { + let mut snapshot = MetricsSnapshot { + samples: Vec::new(), + _strings: Vec::new(), + }; + + for family in registry.gather() { + let name = family.get_name(); + // These request families have a single `type` label in client-rust. + // Restrict both names and types so the C representation is lossless. + if !matches!( + (name, family.get_field_type()), + ( + "tikv_request_total" + | "tikv_failed_request_total" + | "pd_request_total" + | "pd_failed_request_total", + MetricType::COUNTER + ) | ( + "tikv_request_duration_seconds" + | "tikv_failed_request_duration_seconds" + | "pd_request_duration_seconds" + | "pd_failed_request_duration_seconds", + MetricType::HISTOGRAM + ) + ) { + continue; + } + + match family.get_field_type() { + MetricType::COUNTER => { + for metric in family.get_metric() { + let Some(type_label) = metric_type_label(metric) else { + continue; + }; + snapshot.push( + 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 Some(type_label) = metric_type_label(metric) else { + continue; + }; + 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(); + snapshot.push( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET, + type_label, + bucket.get_cumulative_count() as f64, + upper_bound, + ); + } + if !inf_seen { + snapshot.push( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET, + type_label, + histogram.get_sample_count() as f64, + f64::INFINITY, + ); + } + snapshot.push( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_SUM, + type_label, + histogram.get_sample_sum(), + 0.0, + ); + snapshot.push( + name, + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT, + type_label, + histogram.get_sample_count() as f64, + 0.0, + ); + } + } + _ => {} + } + } + + snapshot +} + +#[no_mangle] +pub extern "C" fn tikv_metrics_snapshot_collect( + out_snapshot: *mut *mut tikv_metrics_snapshot_t, + out_samples: *mut tikv_metric_samples_t, + out_error: *mut *mut tikv_error_t, +) -> tikv_status_t { + if !out_snapshot.is_null() { + unsafe { *out_snapshot = ptr::null_mut() }; + } + if !out_samples.is_null() { + unsafe { + *out_samples = tikv_metric_samples_t { + ptr: ptr::null(), + len: 0, + } + }; + } + if out_snapshot.is_null() || out_samples.is_null() { + clear_error_slot(out_error); + let name = if out_snapshot.is_null() { + "out_snapshot" + } else { + "out_samples" + }; + set_error_slot(out_error, format!("{name} must not be null")); + return tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT; + } + + match run_ffi(out_error, || { + let snapshot = Box::new(collect_metrics_snapshot(prometheus::default_registry())); + unsafe { + *out_samples = tikv_metric_samples_t { + ptr: if snapshot.samples.is_empty() { + ptr::null() + } else { + snapshot.samples.as_ptr() + }, + len: snapshot.samples.len(), + }; + *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_destroy(snapshot: *mut tikv_metrics_snapshot_t) { + if snapshot.is_null() { + return; + } + unsafe { + drop(Box::from_raw(snapshot as *mut MetricsSnapshot)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn snapshot_exports_only_client_request_families() { + let registry = prometheus::Registry::new(); + let counter_names = [ + "tikv_request_total", + "tikv_failed_request_total", + "pd_request_total", + "pd_failed_request_total", + ]; + let histogram_names = [ + "tikv_request_duration_seconds", + "tikv_failed_request_duration_seconds", + "pd_request_duration_seconds", + "pd_failed_request_duration_seconds", + ]; + for name in counter_names { + let counter = + prometheus::IntCounterVec::new(prometheus::Opts::new(name, name), &["type"]) + .unwrap(); + counter.with_label_values(&["raw_get"]).inc_by(3); + registry.register(Box::new(counter)).unwrap(); + } + for name in histogram_names { + let histogram = prometheus::HistogramVec::new( + prometheus::HistogramOpts::new(name, name).buckets(vec![0.005, 0.01]), + &["type"], + ) + .unwrap(); + histogram.with_label_values(&["raw_get"]).observe(0.01); + registry.register(Box::new(histogram)).unwrap(); + } + let unrelated_counter = prometheus::IntCounterVec::new( + prometheus::Opts::new("other_requests_total", "other requests"), + &["type", "store"], + ) + .unwrap(); + for store in ["1", "2"] { + unrelated_counter + .with_label_values(&["raw_get", store]) + .inc(); + } + registry.register(Box::new(unrelated_counter)).unwrap(); + for name in ["other_duration_seconds", "pd_tso_batch_size"] { + let histogram = + prometheus::Histogram::with_opts(prometheus::HistogramOpts::new(name, name)) + .unwrap(); + histogram.observe(1.0); + registry.register(Box::new(histogram)).unwrap(); + } + + let snapshot = collect_metrics_snapshot(®istry); + let mut names = HashSet::new(); + let mut counters = 0; + let mut histogram_counts = 0; + for sample in &snapshot.samples { + names.insert(unsafe { CStr::from_ptr(sample.name) }.to_str().unwrap()); + assert_eq!( + unsafe { CStr::from_ptr(sample.type_label) }.to_bytes(), + b"raw_get" + ); + match sample.kind { + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_COUNTER => { + counters += 1; + assert_eq!(sample.value, 3.0); + } + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_BUCKET => { + assert!(sample.le == 0.005 || sample.le == 0.01 || sample.le == f64::INFINITY); + assert_eq!(sample.value, if sample.le < 0.01 { 0.0 } else { 1.0 }); + } + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_SUM => { + assert_eq!(sample.value, 0.01); + } + tikv_metric_sample_kind_t::TIKV_METRIC_SAMPLE_HISTOGRAM_COUNT => { + histogram_counts += 1; + assert_eq!(sample.value, 1.0); + } + kind => panic!("unexpected sample kind: {kind:?}"), + } + } + assert_eq!(counters, counter_names.len()); + assert_eq!(histogram_counts, histogram_names.len()); + assert_eq!( + names, + counter_names.into_iter().chain(histogram_names).collect() + ); + } + + #[test] + fn snapshot_skips_incompatible_types_and_labels() { + let registry = prometheus::Registry::new(); + let counter = prometheus::IntCounterVec::new( + prometheus::Opts::new("tikv_request_total", "extra labels"), + &["type", "store"], + ) + .unwrap(); + for store in ["1", "2"] { + counter.with_label_values(&["raw_get", store]).inc(); + } + registry.register(Box::new(counter)).unwrap(); + let histogram = prometheus::HistogramVec::new( + prometheus::HistogramOpts::new("tikv_request_duration_seconds", "extra labels"), + &["type", "store"], + ) + .unwrap(); + for store in ["1", "2"] { + histogram + .with_label_values(&["raw_get", store]) + .observe(0.01); + } + registry.register(Box::new(histogram)).unwrap(); + let counter = prometheus::IntCounter::new("pd_request_total", "missing label").unwrap(); + counter.inc(); + registry.register(Box::new(counter)).unwrap(); + let histogram = prometheus::Histogram::with_opts(prometheus::HistogramOpts::new( + "pd_request_duration_seconds", + "missing label", + )) + .unwrap(); + histogram.observe(0.01); + registry.register(Box::new(histogram)).unwrap(); + let gauge = prometheus::IntGaugeVec::new( + prometheus::Opts::new("tikv_failed_request_total", "wrong type"), + &["type"], + ) + .unwrap(); + gauge.with_label_values(&["raw_get"]).set(1); + registry.register(Box::new(gauge)).unwrap(); + let counter = prometheus::IntCounterVec::new( + prometheus::Opts::new("tikv_failed_request_duration_seconds", "wrong type"), + &["type"], + ) + .unwrap(); + counter.with_label_values(&["raw_get"]).inc(); + registry.register(Box::new(counter)).unwrap(); + let histogram = prometheus::HistogramVec::new( + prometheus::HistogramOpts::new("pd_failed_request_duration_seconds", "wrong label"), + &["operation"], + ) + .unwrap(); + histogram.with_label_values(&["raw_get"]).observe(0.01); + registry.register(Box::new(histogram)).unwrap(); + + assert!(collect_metrics_snapshot(®istry).samples.is_empty()); + } +} diff --git a/bindings/c/src/raw.rs b/bindings/c/src/raw.rs new file mode 100644 index 00000000..bc7d9cfb --- /dev/null +++ b/bindings/c/src/raw.rs @@ -0,0 +1,434 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +use super::*; +use crate::request::submit_common; + +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() +} + +pub(crate) 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))) + } + MutationOp::Delete => Ok(OwnedMutation::Delete(key)), + } + }) + .collect() +} + +/// Splits a mixed mutation list into ordered PUT and DELETE runs. +/// +/// Each run contains one operation kind and is submitted with the matching +/// `RawClient` batch method. PUT/DELETE transitions preserve the input order. +/// A repeated key ends the current PUT run so later PUTs for the same key are +/// submitted separately and retain their input-order semantics. This function +/// only builds the runs; it does not submit any requests. +pub(crate) 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) => { + 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); + } + 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))) +} + +async fn submit_put_batch(client: &RawClient, puts: Vec) -> FfiResult<()> { + if puts.is_empty() { + return Ok(()); + } + + client + .batch_put(puts) + .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(()) +} + +#[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_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_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_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, + } +} diff --git a/bindings/c/src/request.rs b/bindings/c/src/request.rs new file mode 100644 index 00000000..a763e38e --- /dev/null +++ b/bindings/c/src/request.rs @@ -0,0 +1,613 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +use super::*; + +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 make_request_state( + kind: RequestKind, + lifecycle: &Arc, + callback: tikv_async_completion_cb, + userdata: *mut c_void, +) -> (Arc, InflightGuard) { + let inflight = InflightGuard::new(lifecycle); + let request = Arc::new(RequestState { + kind, + completion: RefCell::new(CompletionState { + complete: false, + completion_status: None, + error_message: None, + output: None, + output_taken: false, + }), + callback: Some(callback), + userdata, + }); + (request, inflight) +} + +fn finalize_request( + request: Arc, + status: tikv_status_t, + output: Option, + error_message: Option, +) { + let (callback, userdata) = { + let mut completion = request.completion.borrow_mut(); + completion.complete = true; + completion.completion_status = Some(status); + completion.output = output; + completion.error_message = + error_message.and_then(|message| CString::new(message.replace('\0', " ")).ok()); + (request.callback, request.userdata) + }; + + if let Some(callback) = callback { + callback( + Arc::as_ptr(&request) as *mut tikv_async_request_t, + status, + userdata, + ); + } +} + +fn spawn_request( + handle: &Handle, + request: Arc, + inflight: InflightGuard, + future: F, +) -> tokio::task::JoinHandle<()> +where + F: std::future::Future> + Send + 'static, +{ + handle.spawn(async move { + // Keep the inflight registration until finalize_request returns, so + // client drain waits for the callback but not for a retained C handle. + let _inflight = inflight; + let result = AssertUnwindSafe(future).catch_unwind().await; + let (status, output, error_message) = match result { + 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, status, output, error_message); + }) +} + +pub(crate) 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 future = build_future(handle.core.client.clone())?; + let (request, inflight) = + make_request_state(kind, &handle.core.lifecycle, completion, userdata); + + if !handle.core.lifecycle.accepting.load(Ordering::Acquire) { + return Err(FfiError::rejected( + "client is shutting down and rejects new requests", + )); + } + + // Transfer one strong reference to C. The task retains the other reference + // through callback completion; C releases its reference via destroy. + unsafe { + *out_request = Arc::into_raw(Arc::clone(&request)) as *mut tikv_async_request_t; + } + spawn_request(handle.core.runtime.handle(), request, inflight, future); + + Ok(()) +} + +fn take_request_output<'a>( + request: &'a RequestState, + expected: RequestKind, +) -> FfiResult> { + if request.kind != expected { + return Err(FfiError::invalid_argument(format!( + "request does not carry a {} result", + expected.name() + ))); + } + let completion = request.completion.borrow_mut(); + 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_request_error_message( + request: *const tikv_async_request_t, +) -> *const c_char { + let Ok(request) = request_ref(request) else { + return ptr::null(); + }; + let completion = request.completion.borrow(); + 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_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_async_request_destroy(request: *mut tikv_async_request_t) { + if request.is_null() { + return; + } + // SAFETY: each accepted request exports exactly one Arc with into_raw. + // The C contract requires releasing that reference exactly once. + unsafe { + drop(Arc::from_raw(request as *const RequestState)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::mpsc; + + fn lifecycle() -> Arc { + Arc::new(ClientLifecycle { + accepting: AtomicBool::new(true), + inflight: Mutex::new(0), + drain_ready: Condvar::new(), + }) + } + + struct CompletionReport { + request_addr: usize, + status: tikv_status_t, + message: Option, + } + + extern "C" fn destroy_in_callback( + request: *mut tikv_async_request_t, + status: tikv_status_t, + userdata: *mut c_void, + ) { + let sender = unsafe { &*(userdata as *const mpsc::Sender) }; + let message = tikv_async_request_error_message(request); + let message = if message.is_null() { + None + } else { + Some( + unsafe { CStr::from_ptr(message) } + .to_string_lossy() + .into_owned(), + ) + }; + tikv_async_request_destroy(request); + let _ = sender.send(CompletionReport { + request_addr: request as usize, + status, + message, + }); + } + + #[test] + fn accepted_requests_complete_once_and_allow_destruction_in_callback() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let lifecycle = lifecycle(); + let (sender, receiver) = mpsc::channel::(); + let outcomes = [ + ( + boxed_future(async { Ok(RequestOutput::Unit) }), + tikv_status_t::TIKV_STATUS_OK, + None, + ), + ( + boxed_future(async { Err(FfiError::error("operation failed")) }), + tikv_status_t::TIKV_STATUS_ERROR, + Some("operation failed"), + ), + ( + boxed_future(async { panic!("operation panicked") }), + tikv_status_t::TIKV_STATUS_PANIC, + Some("panic across FFI boundary: operation panicked"), + ), + ]; + for (future, expected_status, expected_message) in outcomes { + let (request, inflight) = make_request_state( + RequestKind::Put, + &lifecycle, + destroy_in_callback, + &sender as *const _ as *mut c_void, + ); + let weak_request = Arc::downgrade(&request); + let raw_request = Arc::into_raw(Arc::clone(&request)) as *mut tikv_async_request_t; + let task = spawn_request(runtime.handle(), request, inflight, future); + let report = receiver.recv_timeout(Duration::from_secs(5)).unwrap(); + runtime.block_on(task).unwrap(); + assert_eq!(report.request_addr, raw_request as usize); + assert_eq!(report.status, expected_status); + assert_eq!(report.message.as_deref(), expected_message); + assert!(matches!( + receiver.try_recv(), + Err(mpsc::TryRecvError::Empty) + )); + assert_eq!(Arc::strong_count(&lifecycle), 1); + assert!(weak_request.upgrade().is_none()); + } + } + + struct BlockingCallback { + entered: mpsc::Sender, + release: Mutex>, + } + + extern "C" fn blocking_callback( + _request: *mut tikv_async_request_t, + status: tikv_status_t, + userdata: *mut c_void, + ) { + let callback = unsafe { &*(userdata as *const BlockingCallback) }; + let _ = callback.entered.send(status); + let _ = callback.release.lock().unwrap().recv(); + } + + #[test] + fn shutdown_drain_waits_for_pending_operation_and_callback() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let lifecycle = lifecycle(); + let (entered_sender, entered_receiver) = mpsc::channel(); + let (release_sender, release_receiver) = mpsc::channel(); + let callback = BlockingCallback { + entered: entered_sender, + release: Mutex::new(release_receiver), + }; + let (finish_sender, finish_receiver) = futures::channel::oneshot::channel(); + let (request, inflight) = make_request_state( + RequestKind::Put, + &lifecycle, + blocking_callback, + &callback as *const _ as *mut c_void, + ); + let weak_request = Arc::downgrade(&request); + let raw_request = Arc::into_raw(Arc::clone(&request)); + let task = spawn_request(runtime.handle(), request, inflight, async move { + finish_receiver.await.unwrap(); + Ok(RequestOutput::Unit) + }); + lifecycle.accepting.store(false, Ordering::Release); + let (drain_started_sender, drain_started_receiver) = mpsc::channel(); + let (drained_sender, drained_receiver) = mpsc::channel(); + let drain_lifecycle = Arc::clone(&lifecycle); + let drain_thread = std::thread::spawn(move || { + let _ = drain_started_sender.send(()); + wait_drain(&drain_lifecycle); + let _ = drained_sender.send(()); + }); + drain_started_receiver + .recv_timeout(Duration::from_secs(5)) + .unwrap(); + let before_completion = drained_receiver.recv_timeout(Duration::from_millis(50)); + finish_sender.send(()).unwrap(); + let callback_status = entered_receiver + .recv_timeout(Duration::from_secs(5)) + .unwrap(); + let alive_during_callback = weak_request.upgrade().is_some(); + let during_callback = drained_receiver.recv_timeout(Duration::from_millis(50)); + release_sender.send(()).unwrap(); + drained_receiver + .recv_timeout(Duration::from_secs(5)) + .unwrap(); + drain_thread.join().unwrap(); + runtime.block_on(task).unwrap(); + assert_eq!(before_completion, Err(mpsc::RecvTimeoutError::Timeout)); + assert_eq!(during_callback, Err(mpsc::RecvTimeoutError::Timeout)); + assert_eq!(callback_status, tikv_status_t::TIKV_STATUS_OK); + assert!(alive_during_callback); + assert_eq!(Arc::strong_count(&lifecycle), 1); + // The retained C handle keeps the request alive, but no longer keeps + // the client's inflight registration after the callback returns. + assert!(weak_request.upgrade().is_some()); + tikv_async_request_destroy(raw_request as *mut tikv_async_request_t); + assert!(weak_request.upgrade().is_none()); + } + + extern "C" fn report_completion( + _request: *mut tikv_async_request_t, + status: tikv_status_t, + userdata: *mut c_void, + ) { + let sender = unsafe { &*(userdata as *const mpsc::Sender) }; + let _ = sender.send(status); + } + + extern "C" fn handoff_completion( + _request: *mut tikv_async_request_t, + status: tikv_status_t, + userdata: *mut c_void, + ) { + let callback = unsafe { &*(userdata as *const BlockingCallback) }; + // Hand request access to the receiver, then remain in the callback + // without touching the request again. + let _ = callback.entered.send(status); + let _ = callback.release.lock().unwrap().recv(); + } + + #[test] + fn result_can_be_taken_after_handoff_before_callback_returns() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let lifecycle = lifecycle(); + let (entered_sender, entered_receiver) = mpsc::channel(); + let (release_sender, release_receiver) = mpsc::channel(); + let callback = BlockingCallback { + entered: entered_sender, + release: Mutex::new(release_receiver), + }; + let (request, inflight) = make_request_state( + RequestKind::Get, + &lifecycle, + handoff_completion, + &callback as *const _ as *mut c_void, + ); + let weak_request = Arc::downgrade(&request); + let raw_request = Arc::into_raw(Arc::clone(&request)) as *mut tikv_async_request_t; + let task = spawn_request(runtime.handle(), request, inflight, async { + Ok(RequestOutput::Get(Some(b"value".to_vec()))) + }); + let status = entered_receiver + .recv_timeout(Duration::from_secs(5)) + .unwrap(); + + let mut error = ptr::null_mut(); + let mut out_value = empty_owned_bytes(); + let mut out_found = false; + let result_status = tikv_async_request_take_get_result( + raw_request, + &mut out_value, + &mut out_found, + &mut error, + ); + let no_error = error.is_null(); + let found = out_found; + let value = if out_value.ptr.is_null() { + Vec::new() + } else { + unsafe { slice::from_raw_parts(out_value.ptr, out_value.len) }.to_vec() + }; + tikv_owned_bytes_destroy(out_value); + tikv_error_destroy(error); + let second_status = tikv_async_request_take_get_result( + raw_request, + &mut out_value, + &mut out_found, + &mut error, + ); + tikv_error_destroy(error); + tikv_async_request_destroy(raw_request); + let alive_during_callback = weak_request.upgrade().is_some(); + + // Release the callback before asserting so a failed assertion cannot + // leave a runtime worker blocked during runtime destruction. + release_sender.send(()).unwrap(); + runtime.block_on(task).unwrap(); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert_eq!(result_status, tikv_status_t::TIKV_STATUS_OK); + assert!(no_error); + assert!(found); + assert_eq!(value, b"value"); + assert_eq!(second_status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(alive_during_callback); + assert!(weak_request.upgrade().is_none()); + } + + #[test] + fn c_handle_keeps_result_alive_after_task_and_runtime_finish() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let lifecycle = lifecycle(); + let weak_lifecycle = Arc::downgrade(&lifecycle); + let (sender, receiver) = mpsc::channel::(); + let (request, inflight) = make_request_state( + RequestKind::Get, + &lifecycle, + report_completion, + &sender as *const _ as *mut c_void, + ); + let weak_request = Arc::downgrade(&request); + let raw_request = Arc::into_raw(Arc::clone(&request)) as *mut tikv_async_request_t; + let task = spawn_request(runtime.handle(), request, inflight, async { + Ok(RequestOutput::Get(Some(b"value".to_vec()))) + }); + let status = receiver.recv_timeout(Duration::from_secs(5)).unwrap(); + runtime.block_on(task).unwrap(); + drop(runtime); + drop(lifecycle); + + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(weak_request.upgrade().is_some()); + let mut error = ptr::null_mut(); + let mut out_value = empty_owned_bytes(); + let mut out_found = false; + let result_status = tikv_async_request_take_get_result( + raw_request, + &mut out_value, + &mut out_found, + &mut error, + ); + assert_eq!(result_status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert!(out_found); + assert_eq!( + unsafe { slice::from_raw_parts(out_value.ptr, out_value.len) }, + b"value" + ); + + tikv_async_request_destroy(raw_request); + assert!(weak_request.upgrade().is_none()); + assert!(weak_lifecycle.upgrade().is_none()); + // Extracted buffers have their own ownership, independent of the request. + assert_eq!( + unsafe { slice::from_raw_parts(out_value.ptr, out_value.len) }, + b"value" + ); + tikv_owned_bytes_destroy(out_value); + } +} diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs new file mode 100644 index 00000000..63b00dd9 --- /dev/null +++ b/bindings/c/src/tests.rs @@ -0,0 +1,300 @@ +// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0. + +use super::*; +use crate::client::{build_config, parse_pd_endpoints}; +use crate::raw::{build_write_batch_runs, slice_mutations_to_owned}; + +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 { + Arc::into_raw(Arc::new(RequestState { + kind, + completion: RefCell::new(CompletionState { + complete: true, + completion_status: Some(tikv_status_t::TIKV_STATUS_OK), + error_message: None, + output: Some(output), + output_taken: false, + }), + callback: Some(noop_completion), + userdata: ptr::null_mut(), + })) 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 mut samples = tikv_metric_samples_t { + ptr: ptr::null(), + len: 0, + }; + let status = tikv_metrics_snapshot_collect(&mut snapshot, &mut samples, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_OK); + assert!(error.is_null()); + assert!(!snapshot.is_null()); + assert!(!samples.ptr.is_null()); + assert!(samples.len > 0); + + let mut raw_get_total = None; + let mut raw_get_duration_count = None; + let mut saw_inf_bucket = false; + + let sample_array = unsafe { std::slice::from_raw_parts(samples.ptr, samples.len) }; + for sample in sample_array { + 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 metrics_snapshot_collect_rejects_missing_samples_output() { + let mut error = ptr::null_mut(); + let mut snapshot = std::ptr::NonNull::::dangling().as_ptr(); + let status = tikv_metrics_snapshot_collect(&mut snapshot, ptr::null_mut(), &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(snapshot.is_null()); + assert!(!error.is_null()); + tikv_error_destroy(error); + + let mut samples = tikv_metric_samples_t { + ptr: std::ptr::NonNull::::dangling().as_ptr(), + len: 1, + }; + let status = tikv_metrics_snapshot_collect(ptr::null_mut(), &mut samples, &mut error); + assert_eq!(status, tikv_status_t::TIKV_STATUS_INVALID_ARGUMENT); + assert!(samples.ptr.is_null()); + assert_eq!(samples.len, 0); + assert!(!error.is_null()); + tikv_error_destroy(error); +} + +#[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())), + OwnedMutation::Put(KvPair::new(b"b".to_vec(), b"v2".to_vec())), + 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())), + ]); + + assert_eq!( + runs, + vec![ + WriteBatchRun::Put(vec![ + KvPair::new(b"a".to_vec(), b"v1".to_vec()), + KvPair::new(b"b".to_vec(), b"v2".to_vec()), + ]), + 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())]), + ] + ); +} + +#[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())), + OwnedMutation::Put(KvPair::new(b"b".to_vec(), b"v2".to_vec())), + OwnedMutation::Put(KvPair::new(b"a".to_vec(), b"v3".to_vec())), + ]); + + assert_eq!( + runs, + vec![ + WriteBatchRun::Put(vec![ + KvPair::new(b"a".to_vec(), b"v1".to_vec()), + KvPair::new(b"b".to_vec(), b"v2".to_vec()), + ]), + WriteBatchRun::Put(vec![KvPair::new(b"a".to_vec(), b"v3".to_vec())]), + ] + ); +} + +#[test] +fn ffi_mutations_copy_inputs_and_ignore_delete_value() { + let mut key = b"put-key".to_vec(); + let mut value = b"put-value".to_vec(); + 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(), + }, + }, + tikv_async_raw_mutation_t { + op: TIKV_ASYNC_RAW_MUTATION_OP_DELETE, + key: tikv_slice_t { + ptr: key.as_ptr(), + len: key.len(), + }, + value: tikv_slice_t { + ptr: ptr::null(), + len: usize::MAX, + }, + }, + ]; + let owned = slice_mutations_to_owned(tikv_async_raw_mutations_t { + ptr: mutations.as_ptr(), + len: mutations.len(), + }) + .unwrap(); + key.fill(0); + value.fill(0); + assert_eq!( + owned, + vec![ + OwnedMutation::Put(KvPair::new(b"put-key".to_vec(), b"put-value".to_vec())), + OwnedMutation::Delete(Key::from(b"put-key".to_vec())), + ] + ); +} + +#[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_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); +}