diff --git a/client/benchmark.c b/client/benchmark.c index 7da45d3..cb3338e 100644 --- a/client/benchmark.c +++ b/client/benchmark.c @@ -963,8 +963,8 @@ static void zc_alloc_cb(uint64_t rid, priskv_status status, void *result) zctx->pctx->job->mm->memcpy((void *)region->addr, zctx->value, copy_len); zctx->token = region->token; zctx->pctx->job->last_stage = "SEAL"; - priskv_seal_async(zctx->pctx->client, &zctx->token, false /* pin_on_seal */, (uint64_t)zctx, - zc_seal_cb); + priskv_seal_async(zctx->pctx->client, &zctx->token, false /* pin_on_seal */, 0 /* pin_ttl_ms */, + (uint64_t)zctx, zc_seal_cb); } /* ZeroCopy DROP callback is no longer used (published keys use DELETE semantics) */ diff --git a/client/client.c b/client/client.c index 628b8b2..955655a 100644 --- a/client/client.c +++ b/client/client.c @@ -222,7 +222,7 @@ static void set_handler_base(client_context *ctx, char *args, bool alloc) printf("ALLOC_SET status(%d): %s, addr %p, length %u, token 0x%lx\n", status, priskv_status_str(status), (void *)region.addr, region.length, region.token); memcpy((void *)region.addr, value, (size_t)region.length); - status = priskv_seal(ctx->client, ®ion.token, pin_on_seal); + status = priskv_seal(ctx->client, ®ion.token, pin_on_seal, 0 /* pin_ttl_ms */); if (status != PRISKV_STATUS_OK) { printf("Failed to SEAL, status(%d): %s\n", status, priskv_status_str(status)); return; @@ -275,7 +275,7 @@ static void get_handler_base(client_context *ctx, char *args, bool acquire) priskv_memory_region region = {0}; printf("ACQUIRE key=%s\n", key); /* Do not pin on acquire by default from CLI */ - status = priskv_acquire(ctx->client, key, PRISKV_KEY_MAX_TIMEOUT, false, ®ion); + status = priskv_acquire(ctx->client, key, false, 0 /* pin_ttl_ms */, ®ion); if (status != PRISKV_STATUS_OK) { printf("Failed to GET, status(%d): %s\n", status, priskv_status_str(status)); return; @@ -392,6 +392,7 @@ static void seal_token_handler(client_context *ctx, char *args) uint64_t token = 0; priskv_status status; bool pin_on_seal = false; + uint64_t pin_ttl_ms = 0; /* default: server-side TTL */ tokstr = strtok_r(args, " ", &args); if (!tokstr) { @@ -414,33 +415,36 @@ static void seal_token_handler(client_context *ctx, char *args) if (!strcmp(flag, "PIN")) { pin_on_seal = true; } else if (!strcmp(flag, "TTL")) { - /* TODO(wangyi): TTL passthrough is not implemented; parse and discard for now */ char *ttl = strtok_r(args, " ", &args); if (!ttl || !strlen(ttl)) { printf("%s\n", invalid_args_msg); return; } + errno = 0; + pin_ttl_ms = strtoull(ttl, &str_end, 10); + if (errno > 0 || str_end == ttl || *str_end != '\0') { + printf("%s\n", invalid_args_msg); + return; + } } else { printf("%s\n", invalid_args_msg); return; } } - printf("SEAL token=%" PRIu64 " [PIN=%d]\n", token, pin_on_seal); - /* TODO(wangyi): Support per-command TTL for pin-on-seal (e.g., 'PIN TTL N') - * - Parse TTL value and pass through protocol once pin_ttl_ms is supported. - * - Default to server-side TTL when not provided. - */ - status = priskv_seal(ctx->client, &token, pin_on_seal); + printf("SEAL token=%" PRIu64 " [PIN=%d, TTL(ms)=%" PRIu64 "]\n", token, pin_on_seal, + pin_ttl_ms); + status = priskv_seal(ctx->client, &token, pin_on_seal, pin_ttl_ms); printf("SEAL status(%d): %s\n", status, priskv_status_str(status)); } static void acquire_only_handler(client_context *ctx, char *args) { - char *key; + char *key, *str_end; priskv_status status; priskv_memory_region region = {0}; bool pin_on_acquire = false; + uint64_t pin_ttl_ms = 0; /* default: server-side TTL */ key = strtok_r(args, " ", &args); if (!key) { @@ -454,12 +458,17 @@ static void acquire_only_handler(client_context *ctx, char *args) if (!strcmp(flag, "PIN")) { pin_on_acquire = true; } else if (!strcmp(flag, "TTL")) { - /* TODO(wangyi): TTL passthrough is not implemented; parse and discard for now */ char *ttl = strtok_r(args, " ", &args); if (!ttl || !strlen(ttl)) { printf("%s\n", invalid_args_msg); return; } + errno = 0; + pin_ttl_ms = strtoull(ttl, &str_end, 10); + if (errno > 0 || str_end == ttl || *str_end != '\0') { + printf("%s\n", invalid_args_msg); + return; + } } else { printf("%s\n", invalid_args_msg); return; @@ -467,11 +476,8 @@ static void acquire_only_handler(client_context *ctx, char *args) } /* Align output field name with CLI flag semantics */ - printf("ACQUIRE key=%s [PIN=%d]\n", key, pin_on_acquire); - /* TODO(wangyi): Support per-command TTL for pin-on-acquire (e.g., 'PIN TTL N') - * - Parse TTL value and pass through protocol once pin_ttl_ms is supported. - */ - status = priskv_acquire(ctx->client, key, PRISKV_KEY_MAX_TIMEOUT, pin_on_acquire, ®ion); + printf("ACQUIRE key=%s [PIN=%d, TTL(ms)=%" PRIu64 "]\n", key, pin_on_acquire, pin_ttl_ms); + status = priskv_acquire(ctx->client, key, pin_on_acquire, pin_ttl_ms, ®ion); printf("ACQUIRE status(%d): %s, addr %p, length %u, token 0x%lx\n", status, priskv_status_str(status), (void *)region.addr, region.length, region.token); if (status == PRISKV_STATUS_OK) { @@ -545,7 +551,7 @@ static void drop_token_handler(client_context *ctx, char *args) return; } } - printf("DROP token=% \n" PRIu64, token); + printf("DROP token=%" PRIu64 "\n", token); status = priskv_drop(ctx->client, &token); printf("DROP status(%d): %s\n", status, priskv_status_str(status)); } diff --git a/client/priskv.h b/client/priskv.h index f4e77cd..5e5592f 100644 --- a/client/priskv.h +++ b/client/priskv.h @@ -156,24 +156,24 @@ int priskv_flush_async(priskv_client *client, const char *regex, uint64_t reques int priskv_alloc_async(priskv_client *client, const char *key, uint32_t alloc_length, uint64_t timeout, uint64_t request_id, priskv_generic_cb cb); -/* Seal memory region alloced for write (by token pointer, reuse key field) */ +/* Seal memory region alloced for write (by token pointer, reuse key field) + * When pin_on_seal is true, pass TTL via pin_ttl_ms (0 uses server default). + */ int priskv_seal_async(priskv_client *client, const uint64_t *token, bool pin_on_seal, - uint64_t request_id, priskv_generic_cb cb); + uint64_t pin_ttl_ms, uint64_t request_id, priskv_generic_cb cb); -/* Acquire memory region for zero copy read */ -int priskv_acquire_async(priskv_client *client, const char *key, uint64_t timeout, - bool pin_on_acquire, uint64_t request_id, priskv_generic_cb cb); +/* Acquire memory region for zero copy read (pin_ttl_ms is PIN TTL in ms; 0 uses server default) */ +int priskv_acquire_async(priskv_client *client, const char *key, bool pin_on_acquire, + uint64_t pin_ttl_ms, uint64_t request_id, priskv_generic_cb cb); /* Release memory region (by token pointer, reuse key field) */ int priskv_release_async(priskv_client *client, const uint64_t *token, bool unpin_on_release, uint64_t request_id, priskv_generic_cb cb); - /* Drop memory region (by token pointer, reuse key field) */ int priskv_drop_async(priskv_client *client, const uint64_t *token, uint64_t request_id, priskv_generic_cb cb); - /* for *KEYS* command */ typedef struct priskv_key { char *key; @@ -221,10 +221,12 @@ uint64_t priskv_capacity(priskv_client *client); int priskv_alloc(priskv_client *client, const char *key, uint32_t alloc_length, uint64_t timeout, priskv_memory_region *region); -int priskv_seal(priskv_client *client, const uint64_t *token, bool pin_on_seal); +int priskv_seal(priskv_client *client, const uint64_t *token, bool pin_on_seal, + uint64_t pin_ttl_ms); + +int priskv_acquire(priskv_client *client, const char *key, bool pin_on_acquire, uint64_t pin_ttl_ms, + priskv_memory_region *region); -int priskv_acquire(priskv_client *client, const char *key, uint64_t timeout, - bool pin_on_acquire, priskv_memory_region *region); int priskv_release(priskv_client *client, const uint64_t *token, bool unpin_on_release); int priskv_drop(priskv_client *client, const uint64_t *token); diff --git a/client/sync.c b/client/sync.c index 1f8c3fe..bf1fd06 100644 --- a/client/sync.c +++ b/client/sync.c @@ -45,8 +45,9 @@ static void priskv_common_sync_cb(uint64_t request_id, priskv_status status, voi priskv_transport_req_sync *req_sync = (priskv_transport_req_sync *)request_id; uint32_t valuelen = result ? *(uint32_t *)result : 0; - priskv_log_debug("priskv_common_sync_cb: callback request_id 0x%lx, status: %s[0x%x], length %d\n", request_id, - priskv_resp_status_str(status), status, valuelen); + priskv_log_debug( + "priskv_common_sync_cb: callback request_id 0x%lx, status: %s[0x%x], length %d\n", + request_id, priskv_resp_status_str(status), status, valuelen); req_sync->status = status; req_sync->valuelen = valuelen; req_sync->done = true; @@ -61,7 +62,8 @@ static inline int priskv_sync_wait(priskv_client *client, bool *done) return 0; } -int priskv_get(priskv_client *client, const char *key, priskv_sgl *sgl, uint16_t nsgl, uint32_t *valuelen) +int priskv_get(priskv_client *client, const char *key, priskv_sgl *sgl, uint16_t nsgl, + uint32_t *valuelen) { priskv_transport_req_sync req_sync = {.status = 0xffff, .done = false}; @@ -72,7 +74,8 @@ int priskv_get(priskv_client *client, const char *key, priskv_sgl *sgl, uint16_t return req_sync.status; } -int priskv_set(priskv_client *client, const char *key, priskv_sgl *sgl, uint16_t nsgl, uint64_t timeout) +int priskv_set(priskv_client *client, const char *key, priskv_sgl *sgl, uint16_t nsgl, + uint64_t timeout) { priskv_transport_req_sync req_sync = {.status = 0xffff, .done = false}; @@ -177,21 +180,22 @@ int priskv_alloc(priskv_client *client, const char *key, uint32_t alloc_length, return req_sync.status; } -int priskv_seal(priskv_client *client, const uint64_t *token, bool pin_on_seal) +int priskv_seal(priskv_client *client, const uint64_t *token, bool pin_on_seal, uint64_t pin_ttl_ms) { priskv_transport_req_sync req_sync = {.status = 0xffff, .done = false}; - priskv_seal_async(client, token, pin_on_seal, (uint64_t)&req_sync, + priskv_seal_async(client, token, pin_on_seal, pin_ttl_ms, (uint64_t)&req_sync, priskv_common_sync_cb); priskv_sync_wait(client, &req_sync.done); return req_sync.status; } -int priskv_acquire(priskv_client *client, const char *key, uint64_t timeout, - bool pin_on_acquire, priskv_memory_region *region) +int priskv_acquire(priskv_client *client, const char *key, bool pin_on_acquire, uint64_t pin_ttl_ms, + priskv_memory_region *region) { priskv_transport_zero_copy_req_sync req_sync = {.status = 0xffff, .done = false}; - priskv_acquire_async(client, key, timeout, pin_on_acquire, (uint64_t)&req_sync, + /* pin_ttl_ms is per-request PIN TTL in ms; 0 uses server default */ + priskv_acquire_async(client, key, pin_on_acquire, pin_ttl_ms, (uint64_t)&req_sync, priskv_zero_copy_req_sync_cb); priskv_sync_wait(client, &req_sync.done); @@ -232,8 +236,8 @@ static void priskv_keys_sync_cb(uint64_t request_id, priskv_status status, void { priskv_transport_keys_sync *keys_req_sync = (priskv_transport_keys_sync *)request_id; - priskv_log_debug("priskv_keys_sync_cb: callback request_id 0x%lx, status: %s[0x%x]\n", request_id, - priskv_resp_status_str(status), status); + priskv_log_debug("priskv_keys_sync_cb: callback request_id 0x%lx, status: %s[0x%x]\n", + request_id, priskv_resp_status_str(status), status); keys_req_sync->status = status; keys_req_sync->done = true; diff --git a/client/transport/rdma.c b/client/transport/rdma.c index d13c2b9..958e4b0 100644 --- a/client/transport/rdma.c +++ b/client/transport/rdma.c @@ -1152,6 +1152,7 @@ priskv_rdma_req_new(priskv_client *client, priskv_transport_conn *conn, uint64_t rdma_req->main_wq = client->wq; rdma_req->cmd = cmd; rdma_req->timeout = timeout; + rdma_req->pin_ttl_ms = 0; /* default: not set */ rdma_req->key = strdup(key); rdma_req->keylen = keylen; rdma_req->request_id = request_id; @@ -1256,6 +1257,7 @@ static int priskv_rdma_req_send(void *arg) req->command = htobe16(rdma_req->cmd); req->nsgl = htobe16(rdma_req->nsgl); req->timeout = htobe64(rdma_req->timeout); + req->pin_ttl_ms = htobe64(rdma_req->pin_ttl_ms); req->flags = htobe32(rdma_req->req_flags); req->key_length = htobe16(rdma_req->keylen); diff --git a/client/transport/transport.c b/client/transport/transport.c index 3238d37..daad80b 100644 --- a/client/transport/transport.c +++ b/client/transport/transport.c @@ -287,8 +287,8 @@ static inline priskv_transport_conn *priskv_select_conn(priskv_client *client) static void priskv_send_command(priskv_client *client, uint64_t request_id, const char *key, uint32_t alloc_length, priskv_sgl *sgl, uint16_t nsgl, - uint64_t timeout, priskv_req_command cmd, uint32_t req_flags, - priskv_generic_cb cb) + uint64_t key_expiry_timeout, uint64_t pin_ttl_ms, + priskv_req_command cmd, uint32_t req_flags, priskv_generic_cb cb) { priskv_transport_conn *conn = priskv_select_conn(client); priskv_connect_param *param = &conn->param; @@ -326,12 +326,13 @@ static void priskv_send_command(priskv_client *client, uint64_t request_id, cons } req = client->ops->new_req(client, conn, request_id, key_ptr, keylen, alloc_length, sgl, nsgl, - timeout, cmd, cb); + key_expiry_timeout, cmd, cb); if (!req) { cb(request_id, PRISKV_STATUS_NO_MEM, NULL); return; } req->req_flags = req_flags; + req->pin_ttl_ms = pin_ttl_ms; client->ops->submit_req(req); } @@ -343,8 +344,8 @@ int priskv_get_async(priskv_client *client, const char *key, priskv_sgl *sgl, ui return 0; } - priskv_send_command(client, request_id, key, 0 /* alloc_length */, sgl, nsgl, 0, - PRISKV_COMMAND_GET, 0, cb); + priskv_send_command(client, request_id, key, 0 /* alloc_length */, sgl, nsgl, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_GET, 0, cb); return 0; } @@ -356,100 +357,108 @@ int priskv_set_async(priskv_client *client, const char *key, priskv_sgl *sgl, ui return 0; } - priskv_send_command(client, request_id, key, 0 /* alloc_length */, sgl, nsgl, timeout, - PRISKV_COMMAND_SET, 0, cb); + priskv_send_command(client, request_id, key, 0 /* alloc_length */, sgl, nsgl, + timeout /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_SET, 0, + cb); return 0; } int priskv_test_async(priskv_client *client, const char *key, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_TEST, 0, cb); + priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_TEST, 0, cb); return 0; } int priskv_delete_async(priskv_client *client, const char *key, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_DELETE, 0, cb); + priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_DELETE, 0, + cb); return 0; } int priskv_expire_async(priskv_client *client, const char *key, uint64_t timeout, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, timeout, - PRISKV_COMMAND_EXPIRE, 0, cb); + priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, + timeout /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_EXPIRE, + 0, cb); return 0; } int priskv_keys_async(priskv_client *client, const char *regex, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, regex, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_KEYS, 0, cb); + priskv_send_command(client, request_id, regex, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_KEYS, 0, cb); return 0; } int priskv_nrkeys_async(priskv_client *client, const char *regex, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, regex, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_NRKEYS, 0, cb); + priskv_send_command(client, request_id, regex, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_NRKEYS, 0, + cb); return 0; } int priskv_flush_async(priskv_client *client, const char *regex, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, regex, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_FLUSH, 0, cb); + priskv_send_command(client, request_id, regex, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_FLUSH, 0, + cb); return 0; } int priskv_alloc_async(priskv_client *client, const char *key, uint32_t alloc_length, uint64_t timeout, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, key, alloc_length, NULL, 0, timeout, - PRISKV_COMMAND_ALLOC, 0, cb); + priskv_send_command(client, request_id, key, alloc_length, NULL, 0, + timeout /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_ALLOC, + 0, cb); return 0; } int priskv_seal_async(priskv_client *client, const uint64_t *token, bool pin_on_seal, - uint64_t request_id, priskv_generic_cb cb) + uint64_t pin_ttl_ms, uint64_t request_id, priskv_generic_cb cb) { uint32_t flags = pin_on_seal ? PRISKV_REQ_FLAG_PIN_ON_SEAL : 0; - priskv_send_command(client, request_id, (const char *)token, 0 /* alloc_length */, NULL, 0, 0, + priskv_send_command(client, request_id, (const char *)token, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, pin_ttl_ms /* pin_ttl_ms */, PRISKV_COMMAND_SEAL, flags, cb); return 0; } -int priskv_acquire_async(priskv_client *client, const char *key, uint64_t timeout, - bool pin_on_acquire, uint64_t request_id, priskv_generic_cb cb) +int priskv_acquire_async(priskv_client *client, const char *key, bool pin_on_acquire, + uint64_t pin_ttl_ms, uint64_t request_id, priskv_generic_cb cb) { uint32_t flags = pin_on_acquire ? PRISKV_REQ_FLAG_PIN_ON_ACQUIRE : 0; - priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, timeout, + priskv_send_command(client, request_id, key, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, pin_ttl_ms /* pin_ttl_ms */, PRISKV_COMMAND_ACQUIRE, flags, cb); return 0; } int priskv_release_async(priskv_client *client, const uint64_t *token, bool unpin_on_release, - uint64_t request_id, priskv_generic_cb cb) + uint64_t request_id, priskv_generic_cb cb) { uint32_t flags = unpin_on_release ? PRISKV_REQ_FLAG_UNPIN_ON_RELEASE : 0; - priskv_send_command(client, request_id, (const char *)token, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_RELEASE, flags, cb); + priskv_send_command(client, request_id, (const char *)token, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_RELEASE, + flags, cb); return 0; } - int priskv_drop_async(priskv_client *client, const uint64_t *token, uint64_t request_id, priskv_generic_cb cb) { - priskv_send_command(client, request_id, (const char *)token, 0 /* alloc_length */, NULL, 0, 0, - PRISKV_COMMAND_DROP, 0, cb); + priskv_send_command(client, request_id, (const char *)token, 0 /* alloc_length */, NULL, 0, + 0 /* key_expiry_timeout */, 0 /* pin_ttl_ms */, PRISKV_COMMAND_DROP, 0, cb); return 0; } diff --git a/client/transport/transport.h b/client/transport/transport.h index 9d72917..5d70215 100644 --- a/client/transport/transport.h +++ b/client/transport/transport.h @@ -102,6 +102,7 @@ typedef struct priskv_transport_req { uint16_t keylen; uint32_t alloc_length; uint64_t timeout; + uint64_t pin_ttl_ms; /* Per-request PIN TTL in ms; used only when PIN flags are set */ priskv_req_command cmd; uint32_t req_flags; /* Serialize request behavior flags into priskv_request::flags */ void (*cb)(struct priskv_transport_req *req); diff --git a/client/transport/ucx.c b/client/transport/ucx.c index 7483e28..05aa05f 100644 --- a/client/transport/ucx.c +++ b/client/transport/ucx.c @@ -1235,6 +1235,7 @@ priskv_ucx_req_new(priskv_client *client, priskv_transport_conn *conn, uint64_t ucx_req->main_wq = client->wq; ucx_req->cmd = cmd; ucx_req->timeout = timeout; + ucx_req->pin_ttl_ms = 0; /* default: not set */ /* For SEAL/RELEASE/DROP, the key is not a string but a binary token with fixed length keylen */ ucx_req->keylen = keylen; ucx_req->key = malloc(keylen); @@ -1365,6 +1366,7 @@ static int priskv_ucx_send_req(void *arg) req->command = htobe16(ucx_req->cmd); req->nsgl = htobe16(ucx_req->nsgl); req->timeout = htobe64(ucx_req->timeout); + req->pin_ttl_ms = htobe64(ucx_req->pin_ttl_ms); req->flags = htobe32(ucx_req->req_flags); req->key_length = htobe16(ucx_req->keylen); req->alloc_length = htobe32(ucx_req->alloc_length); diff --git a/cluster/client/client.c b/cluster/client/client.c index c253474..53891ea 100644 --- a/cluster/client/client.c +++ b/cluster/client/client.c @@ -107,6 +107,9 @@ struct priskvClusterRequest { const char *key; priskvClusterSGL *cluster_sgl; /* Cluster communication SGL: stores base addresses of communication memory blocks (used for retry) */ uint64_t timeout; + uint64_t pin_ttl_ms; /* PIN TTL in ms; 0 means server default */ + bool pin_key; /* Whether to PIN the key (interpreted by request type) */ + bool unpin_key; /* Whether to UNPIN the key (interpreted by request type) */ uint32_t alloc_length; uint64_t token; /* Token used for SEAL/RELEASE */ priskvClusterNode *node; @@ -242,10 +245,14 @@ static inline int priskvClusterNodeOpen(priskvClusterClient *client, priskvClust static inline void priskvClusterNodeClose(priskvClusterClient *client, priskvClusterNode *node) { - priskv_close(node->client); - node->client = NULL; - free(node->addr); - node->addr = NULL; + if (node->client) { + priskv_close(node->client); + node->client = NULL; + } + if (node->addr) { + free(node->addr); + node->addr = NULL; + } node->port = 0; } @@ -391,6 +398,7 @@ static priskvClusterRequest * priskvClusterRequestNewBase(priskvClusterNode *node, priskvClusterSGL *sgl, uint16_t nsgl, priskvClusterCallback cb, priskvClusterZeroCopyCallback zero_copy_cb, void *cbarg, RequestType type, const char *key, uint64_t timeout, + bool pin_key, uint64_t pin_ttl_ms, bool unpin_key, uint32_t alloc_length, uint64_t token, priskvClusterClient *client) { priskvClusterRequest *req = malloc(sizeof(priskvClusterRequest)); @@ -404,6 +412,9 @@ priskvClusterRequestNewBase(priskvClusterNode *node, priskvClusterSGL *sgl, uint req->key = key; req->cluster_sgl = malloc(sizeof(priskvClusterSGL) * nsgl); req->timeout = timeout; + req->pin_key = pin_key; + req->pin_ttl_ms = pin_ttl_ms; + req->unpin_key = unpin_key; req->alloc_length = alloc_length; req->node = node; req->token = token; @@ -429,16 +440,21 @@ static priskvClusterRequest *priskvClusterRequestNew(priskvClusterNode *node, pr void *cbarg, RequestType type, const char *key, uint64_t timeout, priskvClusterClient *client) { - return priskvClusterRequestNewBase(node, sgl, nsgl, cb, NULL, cbarg, type, key, timeout, 0, 0, - client); + /* For ACQUIRE/SEAL, interpret 'timeout' from callers as pin_ttl_ms to keep API stable */ + return priskvClusterRequestNewBase(node, sgl, nsgl, cb, NULL, cbarg, type, key, timeout, + false /*pin_key*/, 0 /*pin_ttl_ms*/, false /* unpin_key */, + 0, 0, client); } -static priskvClusterRequest *priskvClusterZeroCopyRequestNew( - priskvClusterNode *node, priskvClusterZeroCopyCallback cb, void *cbarg, RequestType type, - const char *key, uint32_t alloc_length, uint64_t timeout, priskvClusterClient *client) +static priskvClusterRequest * +priskvClusterZeroCopyRequestNew(priskvClusterNode *node, priskvClusterZeroCopyCallback cb, + void *cbarg, RequestType type, const char *key, + uint32_t alloc_length, uint64_t timeout, bool pin_key, + uint64_t pin_ttl_ms, bool unpin_key, priskvClusterClient *client) { - return priskvClusterRequestNewBase(node, NULL, 0, NULL, cb, cbarg, type, key, timeout, - alloc_length, 0 /*token*/, client); + /* Zero-copy requests never set unpin_key at construction */ + return priskvClusterRequestNewBase(node, NULL, 0, NULL, cb, cbarg, type, key, timeout, pin_key, + pin_ttl_ms, unpin_key, alloc_length, 0 /*token*/, client); } static void priskvClusterRequestFree(priskvClusterRequest *req) @@ -551,7 +567,8 @@ priskvClusterRequest *priskvClusterGetRequest(priskvClusterClient *client, const priskvClusterRequest *priskvClusterGetZeroCopyRequest(priskvClusterClient *client, const char *key, priskvClusterZeroCopyCallback cb, void *cbarg, uint32_t alloc_length, uint64_t timeout, - RequestType type) + bool pin_key, uint64_t pin_ttl_ms, + bool unpin_key, RequestType type) { priskvClusterNode *node = priskvClusterGetNode(client, key); if (!node) { @@ -559,8 +576,9 @@ priskvClusterRequest *priskvClusterGetZeroCopyRequest(priskvClusterClient *clien return NULL; } - priskvClusterRequest *req = - priskvClusterZeroCopyRequestNew(node, cb, cbarg, type, key, alloc_length, timeout, client); + /* Use explicit pin_ttl_ms for ACQUIRE/SEAL; ALLOC ignores it */ + priskvClusterRequest *req = priskvClusterZeroCopyRequestNew( + node, cb, cbarg, type, key, alloc_length, timeout, pin_key, pin_ttl_ms, unpin_key, client); return req; } @@ -592,15 +610,18 @@ int priskvClusterSubmitRequest(priskvClusterRequest *req) (uint64_t)req, priskvClusterZeroCopyRequestCallback); break; case SEAL: - priskv_seal_async(req->node->client, &req->token, false /* pin_on_seal */, (uint64_t)req, + priskv_seal_async(req->node->client, &req->token, req->pin_key /* pin_on_seal */, + req->pin_ttl_ms /* pin_ttl_ms */, (uint64_t)req, priskvClusterRequestCallback); break; case ACQUIRE: - priskv_acquire_async(req->node->client, req->key, req->timeout, false /* pin_on_acquire */, (uint64_t)req, + priskv_acquire_async(req->node->client, req->key, req->pin_ttl_ms, + req->pin_key /* pin_on_acquire */, (uint64_t)req, priskvClusterZeroCopyRequestCallback); break; case RELEASE: - priskv_release_async(req->node->client, &req->token, false /* unpin_on_release */, (uint64_t)req, + priskv_release_async(req->node->client, &req->token, + req->unpin_key /* unpin_on_release */, (uint64_t)req, priskvClusterRequestCallback); break; case DROP: @@ -904,8 +925,9 @@ int priskvClusterAsyncDelete(priskvClusterClient *client, const char *key, prisk int priskvClusterAsyncAlloc(priskvClusterClient *client, const char *key, uint64_t alloc_length, uint64_t timeout, priskvClusterZeroCopyCallback cb, void *cbarg) { - priskvClusterRequest *req = - priskvClusterGetZeroCopyRequest(client, key, cb, cbarg, alloc_length, timeout, ALLOC); + priskvClusterRequest *req = priskvClusterGetZeroCopyRequest( + client, key, cb, cbarg, alloc_length, timeout, false /* pin_key */, 0 /* pin_ttl_ms */, + false /* unpin_key */, ALLOC); if (req == NULL) return -1; @@ -913,11 +935,14 @@ int priskvClusterAsyncAlloc(priskvClusterClient *client, const char *key, uint64 } int priskvClusterAsyncSeal(priskvClusterClient *client, const char *key, const uint64_t *token, - priskvClusterCallback cb, void *cbarg) + bool pin_on_seal, uint64_t pin_ttl_ms, priskvClusterCallback cb, + void *cbarg) { - priskvClusterRequest *req = priskvClusterGetRequest(client, key, NULL, 0, cb, cbarg, 0, SEAL); + priskvClusterRequest *req = priskvClusterGetRequest( + client, key, NULL, 0, cb, cbarg, pin_ttl_ms /* use timeout as pin_ttl_ms */, SEAL); if (req) { req->token = token ? *token : 0; + req->pin_key = pin_on_seal; } if (req == NULL) return -1; @@ -925,11 +950,12 @@ int priskvClusterAsyncSeal(priskvClusterClient *client, const char *key, const u return priskvClusterSubmitRequest(req); } -int priskvClusterAsyncAcquire(priskvClusterClient *client, const char *key, uint64_t timeout, - priskvClusterZeroCopyCallback cb, void *cbarg) +int priskvClusterAsyncAcquire(priskvClusterClient *client, const char *key, bool pin_on_acquire, + uint64_t pin_ttl_ms, priskvClusterZeroCopyCallback cb, void *cbarg) { priskvClusterRequest *req = priskvClusterGetZeroCopyRequest( - client, key, cb, cbarg, 0 /* aloc_length */, timeout, ACQUIRE); + client, key, cb, cbarg, 0 /* alloc_length */, 0 /* timeout */, pin_on_acquire /* pin_key */, + pin_ttl_ms, false /* unpin_key */, ACQUIRE); if (req == NULL) return -1; @@ -937,12 +963,13 @@ int priskvClusterAsyncAcquire(priskvClusterClient *client, const char *key, uint } int priskvClusterAsyncRelease(priskvClusterClient *client, const char *key, const uint64_t *token, - priskvClusterCallback cb, void *cbarg) + bool unpin_key, priskvClusterCallback cb, void *cbarg) { priskvClusterRequest *req = priskvClusterGetRequest(client, key, NULL, 0, cb, cbarg, 0, RELEASE); if (req) { req->token = token ? *token : 0; + req->unpin_key = unpin_key; } if (req == NULL) return -1; @@ -964,7 +991,6 @@ int priskvClusterAsyncDrop(priskvClusterClient *client, const char *key, const u return priskvClusterSubmitRequest(req); } - priskvClusterStatus priskvClusterGet(priskvClusterClient *client, const char *key, priskvClusterSGL *sgl, uint16_t nsgl, uint32_t *value_len) { @@ -1056,20 +1082,20 @@ int priskvClusterAllocRegion(priskvClusterClient *client, const char *key, uint3 } priskvClusterStatus priskvClusterSeal(priskvClusterClient *client, const char *key, - const uint64_t *token, bool pin_on_seal) + const uint64_t *token, bool pin_on_seal, uint64_t pin_ttl_ms) { priskvClusterNode *node = priskvClusterGetNode(client, key); if (!node) { return PRISKV_CLUSTER_STATUS_NO_SUCH_KEY; } - priskv_status status = priskv_seal(node->client, token, pin_on_seal); + priskv_status status = priskv_seal(node->client, token, pin_on_seal, pin_ttl_ms); return priskvClusterStatusFromPriskvStatus(status); } priskvClusterStatus priskvClusterAcquire(priskvClusterClient *client, const char *key, - uint64_t timeout, bool pin_on_acquire, uint64_t *addr, - uint32_t *valuelen) + uint64_t timeout, bool pin_on_acquire, uint64_t pin_ttl_ms, + uint64_t *addr, uint32_t *valuelen) { priskvClusterNode *node = priskvClusterGetNode(client, key); if (!node) { @@ -1077,7 +1103,7 @@ priskvClusterStatus priskvClusterAcquire(priskvClusterClient *client, const char } priskv_memory_region region = {0}; - priskv_status status = priskv_acquire(node->client, key, timeout, pin_on_acquire, ®ion); + priskv_status status = priskv_acquire(node->client, key, pin_on_acquire, pin_ttl_ms, ®ion); if (status == PRISKV_STATUS_OK) { if (addr) { *addr = region.addr; @@ -1091,14 +1117,16 @@ priskvClusterStatus priskvClusterAcquire(priskvClusterClient *client, const char } int priskvClusterAcquireRegion(priskvClusterClient *client, const char *key, uint64_t timeout, - bool pin_on_acquire, priskv_memory_region *region) + bool pin_on_acquire, uint64_t pin_ttl_ms, + priskv_memory_region *region) { priskvClusterNode *node = priskvClusterGetNode(client, key); if (!node) { return PRISKV_CLUSTER_STATUS_NO_SUCH_KEY; } - return priskv_acquire(node->client, key, timeout, pin_on_acquire, region); + /* timeout is kept for API compatibility; pin_ttl_ms controls PIN TTL (0 uses server default) */ + return priskv_acquire(node->client, key, pin_on_acquire, pin_ttl_ms, region); } priskvClusterStatus priskvClusterRelease(priskvClusterClient *client, const char *key, diff --git a/cluster/client/client.h b/cluster/client/client.h index 8167672..3797de7 100644 --- a/cluster/client/client.h +++ b/cluster/client/client.h @@ -87,11 +87,12 @@ int priskvClusterAsyncDelete(priskvClusterClient *client, const char *key, prisk int priskvClusterAsyncAlloc(priskvClusterClient *client, const char *key, uint64_t alloc_length, uint64_t timeout, priskvClusterZeroCopyCallback cb, void *cbarg); int priskvClusterAsyncSeal(priskvClusterClient *client, const char *key, const uint64_t *token, - priskvClusterCallback cb, void *cbarg); -int priskvClusterAsyncAcquire(priskvClusterClient *client, const char *key, uint64_t timeout, - priskvClusterZeroCopyCallback cb, void *cbarg); + bool pin_on_seal, uint64_t pin_ttl_ms, priskvClusterCallback cb, + void *cbarg); +int priskvClusterAsyncAcquire(priskvClusterClient *client, const char *key, bool pin_on_acquire, + uint64_t pin_ttl_ms, priskvClusterZeroCopyCallback cb, void *cbarg); int priskvClusterAsyncRelease(priskvClusterClient *client, const char *key, const uint64_t *token, - priskvClusterCallback cb, void *cbarg); + bool unpin_key, priskvClusterCallback cb, void *cbarg); int priskvClusterAsyncDrop(priskvClusterClient *client, const char *key, const uint64_t *token, priskvClusterCallback cb, void *cbarg); /* sync APIs */ @@ -102,9 +103,9 @@ priskvClusterStatus priskvClusterSet(priskvClusterClient *client, const char *ke priskvClusterStatus priskvClusterAlloc(priskvClusterClient *client, const char *key, uint64_t alloc_length, uint64_t timeout, uint64_t *addr); priskvClusterStatus priskvClusterSeal(priskvClusterClient *client, const char *key, - const uint64_t *token, bool pin_on_seal); + const uint64_t *token, bool pin_on_seal, uint64_t pin_ttl_ms); priskvClusterStatus priskvClusterAcquire(priskvClusterClient *client, const char *key, - uint64_t timeout, bool pin_on_acquire, + uint64_t timeout, bool pin_on_acquire, uint64_t pin_ttl_ms, uint64_t *addr_offset, uint32_t *valuelen); priskvClusterStatus priskvClusterRelease(priskvClusterClient *client, const char *key, const uint64_t *token, bool unpin_on_release); @@ -120,4 +121,5 @@ priskvClusterStatus priskvClusterStatusFromPriskvStatus(priskv_status status); int priskvClusterAllocRegion(priskvClusterClient *client, const char *key, uint32_t alloc_length, uint64_t timeout, priskv_memory_region *region); int priskvClusterAcquireRegion(priskvClusterClient *client, const char *key, uint64_t timeout, - bool pin_on_acquire, priskv_memory_region *region); + bool pin_on_acquire, uint64_t pin_ttl_ms, + priskv_memory_region *region); diff --git a/include/priskv-protocol.h b/include/priskv-protocol.h index 945e5a2..3812517 100644 --- a/include/priskv-protocol.h +++ b/include/priskv-protocol.h @@ -121,6 +121,7 @@ typedef struct priskv_request_runtime { typedef struct priskv_request { uint64_t request_id; uint64_t timeout; /* in ms */ + uint64_t pin_ttl_ms; /* in ms: 0 uses server default TTL; >0 uses explicit TTL */ uint16_t command; /* priskv_req_command */ uint32_t flags; /* request behavior flags (big-endian on wire) */ uint16_t nsgl; /* how many SGL contains following */ diff --git a/pypriskv/priskv/priskv_client.py b/pypriskv/priskv/priskv_client.py index 58b7ead..bd0157a 100644 --- a/pypriskv/priskv/priskv_client.py +++ b/pypriskv/priskv/priskv_client.py @@ -76,15 +76,18 @@ def alloc(self, def seal(self, key: str, region: client.MemoryRegion, - pin_on_seal: bool = False) -> int: - return client.seal(self.conn, key, region, pin_on_seal) + pin_on_seal: bool = False, + pin_ttl_ms: int = 0) -> int: + """Seal the previously allocated region. When pin_on_seal is True, pin_ttl_ms specifies the PIN TTL in milliseconds; 0 means server default.""" + return client.seal(self.conn, key, region, pin_on_seal, pin_ttl_ms) def acquire(self, key: str, - timeout: int = client.PRISKV_KEY_MAX_TIMEOUT, + pin_ttl_ms: int = 0, pin_on_acquire: bool = False, ) -> Tuple[int, client.MemoryRegion]: - status, region = client.acquire(self.conn, key, timeout, pin_on_acquire) + """Acquire zero-copy read region. When pin_on_acquire is True, pin_ttl_ms specifies the PIN TTL in milliseconds; 0 means server default.""" + status, region = client.acquire(self.conn, key, pin_on_acquire, pin_ttl_ms) return status, region def release(self, diff --git a/pypriskv/pybind.cpp b/pypriskv/pybind.cpp index ee45724..f0a6576 100644 --- a/pypriskv/pybind.cpp +++ b/pypriskv/pybind.cpp @@ -139,13 +139,13 @@ std::tuple priskv_alloc_wrapper(uintptr_t client, std::string key } std::tuple priskv_acquire_wrapper(uintptr_t client, std::string key, - uint64_t timeout, bool pin_on_acquire) + uint64_t pin_ttl_ms, bool pin_on_acquire) { uint64_t addr_offset = 0; uint32_t value_length = 0; - int ret = priskvClusterAcquire((priskvClusterClient *)client, key.c_str(), timeout, - pin_on_acquire, &addr_offset, &value_length); + int ret = priskvClusterAcquire((priskvClusterClient *)client, key.c_str(), 0 /* timeout */, + pin_on_acquire, pin_ttl_ms, &addr_offset, &value_length); return {ret, addr_offset, value_length}; } @@ -315,22 +315,30 @@ PYBIND11_MODULE(_priskv_client, m) m.def( "seal", - [](uintptr_t client, std::string key, const priskv_memory_region ®ion, bool pin_on_seal) { - return (int)priskvClusterSeal((priskvClusterClient *)client, key.c_str(), ®ion.token, pin_on_seal); + [](uintptr_t client, std::string key, const priskv_memory_region ®ion, bool pin_on_seal, + uint64_t pin_ttl_ms) { + return (int)priskvClusterSeal((priskvClusterClient *)client, key.c_str(), ®ion.token, + pin_on_seal, pin_ttl_ms); }, py::arg("client"), py::arg("key"), py::arg("region"), py::arg("pin_on_seal") = false, - "A function to seal memory region of val."); + py::arg("pin_ttl_ms") = 0, + "A function to seal memory region of val with optional PIN TTL (ms). 0 uses server " + "default."); m.def( "acquire", - [](uintptr_t client, std::string key, uint64_t timeout, bool pin_on_acquire) { + [](uintptr_t client, std::string key, bool pin_on_acquire, uint64_t pin_ttl_ms) { priskv_memory_region region {0}; - int ret = priskvClusterAcquireRegion((priskvClusterClient *)client, key.c_str(), - timeout, pin_on_acquire, ®ion); + /* Pass timeout=0 and explicit pin_ttl_ms */ + int ret = + priskvClusterAcquireRegion((priskvClusterClient *)client, key.c_str(), + 0 /* timeout */, pin_on_acquire, pin_ttl_ms, ®ion); return py::make_tuple(ret, region); }, - py::arg("client"), py::arg("key"), py::arg("timeout"), py::arg("pin_on_acquire") = false, - "A function to acquire memory region for read."); + py::arg("client"), py::arg("key"), py::arg("pin_on_acquire") = false, + py::arg("pin_ttl_ms") = 0, + "A function to acquire memory region for read; when PIN is requested, pin_ttl_ms sets TTL " + "in ms (0 uses server default)."); m.def( "release", diff --git a/pypriskv/testing.py b/pypriskv/testing.py index 27d08f7..c05a1ab 100644 --- a/pypriskv/testing.py +++ b/pypriskv/testing.py @@ -26,6 +26,7 @@ import priskv import argparse import ctypes +import time """ Zero-Copy Transport Semantics Assertions (for tests): @@ -536,6 +537,38 @@ def test_unpin_no_such_key(self): assert status in (priskv.PRISKV_STATUS.PRISKV_STATUS_NO_SUCH_TOKEN, priskv.PRISKV_STATUS.PRISKV_STATUS_NO_SUCH_KEY) + def test_pin_ttl_expire_on_acquire(self): + TEST_KEY = "py_pin_ttl_acquire" + SIZE = 256 + # TTL is in milliseconds; sleep slightly longer than 1.5s to avoid clock granularity issues + PIN_TTL_MS = 1500 + SLEEP_SEC = 1.6 + + status, region = self.client.alloc(TEST_KEY, SIZE, 3000) + assert status == 0 + status = self.client.seal(TEST_KEY, region) + assert status == 0 + + # ACQUIRE with PIN（with TTL） + status, acq = self.client.acquire(TEST_KEY, PIN_TTL_MS, True) + assert status == 0, f"acquire(pin ttl) failed: {status}" + + # wait TTL expired + time.sleep(SLEEP_SEC) + + status = self.client.release(TEST_KEY, acq, unpin_on_release=True) + assert status == priskv.PRISKV_STATUS.PRISKV_STATUS_UNPIN_NOT_CLOSED, \ + f"expected UNPIN_NOT_CLOSED after TTL expiry, got {status}" + + status2, acq2 = self.client.acquire(TEST_KEY, 0, False) + assert status2 == priskv.PRISKV_STATUS.PRISKV_STATUS_OK + status2 = self.client.release(TEST_KEY, acq2, unpin_on_release=True) + assert status2 == priskv.PRISKV_STATUS.PRISKV_STATUS_UNPIN_NOT_CLOSED, \ + f"follow-up unpin without pin should be UNPIN_NOT_CLOSED, got {status2}" + + # Cleanup + assert self.client.delete(TEST_KEY) == 0 + # TODO(wangyi): Add PinTTL expiry tests # - Simulate pin-on-acquire/pin-on-seal with a short TTL and verify automatic cleanup # decrements pin_count on the latest version after TTL. @@ -563,6 +596,7 @@ def run_testing(testing): testing.test_pin_on_acquire_and_unpin_on_release() testing.test_unpin_not_closed_on_release() testing.test_unpin_no_such_key() + testing.test_pin_ttl_expire_on_acquire() testing.cleanup() diff --git a/run_unit_test.py b/run_unit_test.py index 2d4845d..611d7f3 100755 --- a/run_unit_test.py +++ b/run_unit_test.py @@ -34,7 +34,8 @@ def priskv_unit_test(parallel: bool = True): "./server/test/test-slab-mt", "./server/test/test-buddy", "./server/test/test-buddy-mt", "./server/test/test-kv", "./server/test/test-kv-mt", "./server/test/test-memory --no-tmpfs", - "./server/test/test-slab", "./server/test/test-transport" + "./server/test/test-slab", "./server/test/test-transport", + "./server/test/test-kv-pin-ttl" ] print("---- PrisKV UNIT TEST ----") diff --git a/server/info.c b/server/info.c index da51356..25e9aa1 100644 --- a/server/info.c +++ b/server/info.c @@ -27,7 +27,6 @@ #include #include #include -#include #include #include #include @@ -82,18 +81,9 @@ void priskv_info_get_kv(void *data) info->expire_kv_bytes = priskv_get_expire_kv_bytes(kv); /* pin/unpin observability */ info->pin_ops = priskv_get_pin_ops(kv); + info->pin_failed_ops = priskv_get_pin_failed_ops(kv); info->unpin_ops = priskv_get_unpin_ops(kv); info->unpin_not_closed = priskv_get_unpin_not_closed(kv); - - /* - * TODO(wangyi): PinTTL observability - * - Once PinTTL manager is introduced, expose additional counters here, e.g.: - * - pin_ttl_active: current number of active PinOperator entries - * - pin_ttl_expired: total number of expired PinOperator entries observed - * - pin_ttl_cleanup_ops: total number of successful cleanup-unpin operations - * - pin_ttl_orphaned: number of expiries that found NO_SUCH_KEY or mismatched state - * - Also consider exporting a configurable default TTL (kv->pin_ttl_ms) for introspection. - */ } void priskv_info_get_connection(void *data) diff --git a/server/jsonobjs.c b/server/jsonobjs.c index 68ee808..4fe9181 100644 --- a/server/jsonobjs.c +++ b/server/jsonobjs.c @@ -68,6 +68,8 @@ PRISKV_DECL_OBJECT_VALUE_FIELD(priskv_kv_info, "expire_kv_bytes", expire_kv_byte required, forced) PRISKV_DECL_OBJECT_VALUE_FIELD(priskv_kv_info, "pin_ops", pin_ops, priskv_uint64, required, forced) +PRISKV_DECL_OBJECT_VALUE_FIELD(priskv_kv_info, "pin_failed_ops", pin_failed_ops, priskv_uint64, + required, forced) PRISKV_DECL_OBJECT_VALUE_FIELD(priskv_kv_info, "unpin_ops", unpin_ops, priskv_uint64, required, forced) PRISKV_DECL_OBJECT_VALUE_FIELD(priskv_kv_info, "unpin_not_closed", unpin_not_closed, priskv_uint64, diff --git a/server/jsonobjs.h b/server/jsonobjs.h index 44c8b30..679b578 100644 --- a/server/jsonobjs.h +++ b/server/jsonobjs.h @@ -65,6 +65,7 @@ typedef struct priskv_kv_info { uint64_t expire_kv_bytes; /* pin/unpin observability */ uint64_t pin_ops; + uint64_t pin_failed_ops; uint64_t unpin_ops; uint64_t unpin_not_closed; /* TODO(wangyi): Extend with PinTTL metrics once available diff --git a/server/kv.c b/server/kv.c index d6d4d6c..0daee68 100644 --- a/server/kv.c +++ b/server/kv.c @@ -101,69 +101,12 @@ typedef struct priskv_kv { void *mf_ctx; struct { uint64_t pin_ops; + uint64_t pin_failed_ops; /* number of failed PIN attempts */ uint64_t unpin_ops; uint64_t unpin_not_closed; } pin_stats; } priskv_kv; -/* - * TODO(wangyi): Implement PinTTL cleanup mechanism - * - * Context: - * - Pin operations (PIN_ON_ACQUIRE, PIN_ON_SEAL) increase pin_count on the latest visible - * version to protect keys from eviction. If a consumer crashes or a request fails, some - * pin operations may never be closed by UNPIN (e.g., RELEASE with UNPIN_ON_RELEASE), leaving - * keys indefinitely pinned. - * - * Goals: - * - Introduce a best-effort TTL-based cleanup for orphaned pins so that keys are eventually - * unpinned when their associated requests are gone. - * - Maintain multi-version correctness: UNPIN targets the latest version even if the pin was - * created before a SEAL publish that replaced the visible version. - * - * Proposed design: - * - PinOperator: a lightweight record created on each effective pin, containing: - * - key (bytes + length) - * - creation timestamp (monotonic clock) - * - ttl_ms (configurable per pin or global default) - * - optional origin (ACQUIRE or SEAL) and a debug request_id for observability - * - PinManager: per-bucket or global manager storing PinOperator entries in an expiry-ordered - * min-heap or timing-wheel to enable O(logN) insert and efficient batch expiry checks. - * - On pin: - * - After pin_count++ on the targeted keynode, create and register a PinOperator. - * - On unpin: - * - Remove the corresponding PinOperator (match by key); then decrement pin_count on latest - * version using priskv_key_unpin_latest semantics. Multiple pins on the same key will - * have multiple PinOperator entries. - * - On seal (version migration): - * - pin_count is already inherited to the new version; PinOperator records keep referencing - * the key (not the keynode pointer), so no migration is required. - * - Scheduling: - * - Reuse expire routine infrastructure (timerfd) to periodically check PinManager and - * perform cleanup for expired entries. Each expired PinOperator triggers - * priskv_key_unpin_latest(kv, old_keynode_of_record) on the latest version by key. - * - Consider sharding the PinManager by hash-bucket index to minimize global contention. - * - Concurrency & locking: - * - PinOperator insert/remove should use per-bucket spinlocks consistent with hash-head - * protection, avoiding deadlocks by keeping lock order (manager lock -> keynode lock). - * - Configuration: - * - Provide a global default TTL (e.g., kv->pin_ttl_ms) and allow per-request override via - * request flags or auxiliary fields in the protocol header (future extension). - * - Observability & safeguards: - * - Export counters: pin_ttl_active, pin_ttl_expired, pin_ttl_cleanup_ops, pin_ttl_orphaned. - * - Cap the maximum number of active PinOperator entries to prevent memory blow-up; when - * exceeding the cap, log warnings and fallback to immediate unpin or refuse new pins. - * - Recovery: - * - PinTTL metadata is best-effort and may be non-persistent. After restart, keys might - * remain pinned by pin_count; scheduled cleanup resumes with new PinOperator records for - * future pins. Persistent logging can be considered if stronger guarantees are needed. - * - * Next steps: - * - Add PinManager data structures and lifecycle APIs. - * - Integrate with pin/unpin paths and expire routine scheduling. - * - Extend info endpoints to expose PinTTL metrics. - */ - static void priskv_lru_access(priskv_key *keynode, bool is_in_list) { priskv_kv *kv = keynode->kv; @@ -290,6 +233,7 @@ void *priskv_new_kv(uint8_t *key_base, uint8_t *value_base, int shm_fd, uint64_t kv->mf_ctx = mf_ctx; kv->pin_stats.pin_ops = 0; + kv->pin_stats.pin_failed_ops = 0; kv->pin_stats.unpin_ops = 0; kv->pin_stats.unpin_not_closed = 0; @@ -394,7 +338,7 @@ static void priskv_keynode_deref(priskv_key *keynode) void priskv_update_valuelen(void *arg, uint32_t valuelen) { priskv_key *keynode = arg; - + pthread_spin_lock(&keynode->lock); keynode->valuelen = valuelen; pthread_spin_unlock(&keynode->lock); @@ -451,25 +395,64 @@ static priskv_key *priskv_find_key(priskv_kv *kv, uint8_t *key, uint16_t keylen, return NULL; } -/* - * Apply a delta to pin_count on the latest visible version of a key. - * - Look up the latest version under the hash-bucket lock. - * - If the key is expired, remove it and return NO_SUCH_KEY. - * - For delta > 0: increment pin_count by delta. - * - For delta < 0: decrement pin_count by |delta| if possible; otherwise return - * PRISKV_RESP_STATUS_UNPIN_NOT_CLOSED without modifying pin_count. - * - Does not update kv->pin_stats; the caller is responsible for stats accounting. +/* Atomically apply a delta to pin_count on the latest visible version. + * - delta > 0: increment and register one TTL entry (ttl_ms == 0 uses default). + * - delta < 0: decrement and unregister one TTL entry. */ static priskv_resp_status priskv_pin_count_delta_latest(priskv_kv *kv, uint8_t *key, - uint16_t keylen, int32_t delta); + uint16_t keylen, int32_t delta, + uint64_t ttl_ms); static void __priskv_del_key(priskv_kv *kv, priskv_key *keynode) { priskv_keynode_deref(keynode); } +/* Increment pin_count and extend node-level TTL. Caller must hold node->lock. */ +static inline void __pin_update_locked(priskv_key *node, uint64_t ttl_ms) +{ + /* Maintain pin_count here to centralize PIN semantics under node->lock. */ + node->pin_count++; + + /* Node-level TTL: choose the later of current vs new (never shorten TTL) */ + if (ttl_ms == 0) + ttl_ms = PRISKV_DEFAULT_PIN_TTL_MS; + + struct timeval now, new_ttl; + gettimeofday(&now, NULL); + new_ttl = now; + priskv_time_add_ms(&new_ttl, ttl_ms); + + /* If TTL already valid, keep the later one; else set to new */ + if (node->pin_ttl.tv_sec >= 0) { + /* priskv_time_elapsed_ms(old, new) > 0 means old < new (new is later); pick new */ + if (priskv_time_elapsed_ms(node->pin_ttl, new_ttl) > 0) { + node->pin_ttl = new_ttl; + } /* else: keep the existing later TTL */ + } else { + node->pin_ttl = new_ttl; + } +} + +/* Decrement pin_count if it's positive. Caller must hold node->lock. */ +static inline void __pin_count_dec_if_positive(priskv_key *node) +{ + if (node->pin_count > 0) { + node->pin_count--; + /* When pin_count drops to zero, clear pin_ttl to invalid + * to reflect that no active pins remain. */ + if (node->pin_count == 0) { + node->pin_ttl.tv_sec = -1; + node->pin_ttl.tv_usec = -1; + } + } +} + +/* Simplified: manage per-key pin_count and TTL directly under bucket lock. + * Caller holds bucket lock, then node lock, to adjust pin_count atomically. */ static priskv_resp_status priskv_pin_count_delta_latest(priskv_kv *kv, uint8_t *key, - uint16_t keylen, int32_t delta) + uint16_t keylen, int32_t delta, + uint64_t ttl_ms) { if (!kv || !key || !keylen) { return PRISKV_RESP_STATUS_SERVER_ERROR; @@ -483,7 +466,8 @@ static priskv_resp_status priskv_pin_count_delta_latest(priskv_kv *kv, uint8_t * gettimeofday(&now, NULL); pthread_spin_lock(&hash_head->lock); - priskv_key *cur; priskv_key *latest = NULL; + priskv_key *cur; + priskv_key *latest = NULL; list_for_each (&hash_head->head, cur, entry) { if (cur->keylen == keylen && memcmp(cur->key, key, keylen) == 0) { latest = cur; @@ -506,24 +490,53 @@ static priskv_resp_status priskv_pin_count_delta_latest(priskv_kv *kv, uint8_t * } priskv_resp_status resp = PRISKV_RESP_STATUS_OK; - pthread_spin_lock(&latest->lock); - if (delta >= 0) { - latest->pin_count += (uint32_t)delta; - } else { - uint32_t need = (uint32_t)(-delta); - if (latest->pin_count >= need) { - latest->pin_count -= need; + if (delta > 0) { + /* PIN (+1): centralized increment + TTL extension */ + pthread_spin_lock(&latest->lock); + __pin_update_locked(latest, ttl_ms); + pthread_spin_unlock(&latest->lock); + } else if (delta < 0) { + /* UNPIN (-1): use pin_count helpers to simplify logic */ + pthread_spin_lock(&latest->lock); + /* If pin TTL has expired, treat as NOT_CLOSED regardless of pin_count. + * This makes UNPIN idempotent after TTL expiration and avoids pretending + * a valid close when the lifetime already ended asynchronously. + */ + bool ttl_expired = false; + if (latest->pin_ttl.tv_sec >= 0) { + int64_t elapsed_ms = priskv_time_elapsed_ms(latest->pin_ttl, now); + if (elapsed_ms > 0) { + ttl_expired = true; + } + } + + if (ttl_expired) { + /* Converge state: clear pin_count and invalidate pin_ttl */ + latest->pin_count = 0; + latest->pin_ttl.tv_sec = -1; + latest->pin_ttl.tv_usec = -1; + resp = PRISKV_RESP_STATUS_UNPIN_NOT_CLOSED; + } else if (latest->pin_count > 0) { + __pin_count_dec_if_positive(latest); } else { - /* Not enough pins closed by the caller */ resp = PRISKV_RESP_STATUS_UNPIN_NOT_CLOSED; } + pthread_spin_unlock(&latest->lock); + + if (resp == PRISKV_RESP_STATUS_UNPIN_NOT_CLOSED && ttl_expired) { + priskv_log_info("KV: UNPIN_NOT_CLOSED due to pin TTL expiry"); + } + } else { + /* delta == 0: no-op */ } - pthread_spin_unlock(&latest->lock); + pthread_spin_unlock(&hash_head->lock); return resp; } +/* no global PinTTL manager; per-key registration/unregistration helpers implemented above */ + int priskv_get_key_for_seal(void *_kv, uint8_t *key, uint16_t keylen, uint8_t **val, uint32_t *valuelen, void **_keynode) { @@ -666,6 +679,8 @@ int priskv_set_key(void *_kv, uint8_t *key, uint16_t keylen, uint8_t **val, uint memcpy(keynode->key, key, keylen); keynode->refcnt = 0; keynode->pin_count = 0; + keynode->pin_ttl.tv_sec = -1; + keynode->pin_ttl.tv_usec = -1; pthread_spin_init(&keynode->lock, 0); priskv_keynode_ref(keynode); @@ -766,6 +781,8 @@ int priskv_alloc_node_private(void *_kv, uint8_t *key, uint16_t keylen, uint8_t } keynode->refcnt = 0; keynode->pin_count = 0; + keynode->pin_ttl.tv_sec = -1; + keynode->pin_ttl.tv_usec = -1; pthread_spin_init(&keynode->lock, 0); priskv_keynode_ref(keynode); /* Initial reference held by the token */ @@ -790,10 +807,9 @@ int priskv_alloc_node_private(void *_kv, uint8_t *key, uint16_t keylen, uint8_t * new node into the hash table. * - Join LRU and clear inprocess so it becomes readable (ACQUIRE/GET). */ -int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish) +static int __priskv_publish_node_with_pin(priskv_kv *kv, priskv_key *keynode, bool pin_on_publish, + uint64_t ttl_ms) { - priskv_kv *kv = _kv; - priskv_key *keynode = (priskv_key *)_keynode; priskv_key *old_keynode = NULL; if (!keynode || keynode->kv != kv) { @@ -816,15 +832,16 @@ int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish) } } - /* Inherit pin_count from old version before making the new one visible. */ + /* Atomic: inherit pin_count under bucket lock. + * Lock order: bucket -> old node -> new node. */ if (old_keynode) { pthread_spin_lock(&old_keynode->lock); - uint32_t old_pins = old_keynode->pin_count; - pthread_spin_unlock(&old_keynode->lock); - pthread_spin_lock(&keynode->lock); - keynode->pin_count = old_pins; + /* inherit pin_count and pin_ttl */ + keynode->pin_count = old_keynode->pin_count; + keynode->pin_ttl = old_keynode->pin_ttl; pthread_spin_unlock(&keynode->lock); + pthread_spin_unlock(&old_keynode->lock); /* Remove old from hash list (visibility) while holding the bucket lock. */ list_del(&old_keynode->entry); @@ -832,18 +849,20 @@ int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish) /* No old version: initialize pin_count to 0 for safety. */ pthread_spin_lock(&keynode->lock); keynode->pin_count = 0; + keynode->pin_ttl.tv_sec = -1; + keynode->pin_ttl.tv_usec = -1; pthread_spin_unlock(&keynode->lock); } - /* Optional: pin on publish within the same critical section to ensure atomicity - * relative to visibility. */ + /* Optional: pin-on-publish within the same critical section (with TTL). */ if (pin_on_publish) { pthread_spin_lock(&keynode->lock); - keynode->pin_count++; + /* Centralized increment + TTL extension */ + __pin_update_locked(keynode, ttl_ms); pthread_spin_unlock(&keynode->lock); } - /* Insert new node into hash list under the same lock to ensure single visible version. */ + /* Make the new node visible in the same critical section. */ list_add_tail(&hash_head->head, &keynode->entry); pthread_spin_unlock(&hash_head->lock); @@ -860,9 +879,17 @@ int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish) return PRISKV_RESP_STATUS_OK; } +int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish, uint64_t ttl_ms) +{ + priskv_kv *kv = _kv; + priskv_key *keynode = (priskv_key *)_keynode; + return __priskv_publish_node_with_pin(kv, keynode, pin_on_publish, ttl_ms); +} + int priskv_publish_node(void *_kv, void *_keynode) { - return priskv_publish_node_with_pin(_kv, _keynode, false); + /* No pin on publish; ttl_ms is irrelevant, pass 0. */ + return priskv_publish_node_with_pin(_kv, _keynode, false, 0); } /* @@ -892,7 +919,6 @@ int priskv_drop_node(void *_kv, void *_keynode) return PRISKV_RESP_STATUS_OK; } - /* Decrement pin_count on the latest version of the key corresponding to keynode. */ int priskv_key_unpin_latest(void *_kv, void *_keynode) { @@ -902,7 +928,7 @@ int priskv_key_unpin_latest(void *_kv, void *_keynode) return PRISKV_RESP_STATUS_SERVER_ERROR; } - priskv_resp_status resp = priskv_pin_count_delta_latest(kv, node->key, node->keylen, -1); + priskv_resp_status resp = priskv_pin_count_delta_latest(kv, node->key, node->keylen, -1, 0); /* Stats: count every UNPIN attempt; record not-closed cases. */ kv->pin_stats.unpin_ops++; @@ -913,7 +939,7 @@ int priskv_key_unpin_latest(void *_kv, void *_keynode) } /* Increment pin_count on the latest visible version of the key corresponding to keynode. */ -int priskv_key_pin_latest(void *_kv, void *_keynode) +int priskv_key_pin_latest(void *_kv, void *_keynode, uint64_t ttl_ms) { priskv_kv *kv = (priskv_kv *)_kv; priskv_key *node = (priskv_key *)_keynode; @@ -921,9 +947,13 @@ int priskv_key_pin_latest(void *_kv, void *_keynode) return PRISKV_RESP_STATUS_SERVER_ERROR; } - priskv_resp_status resp = priskv_pin_count_delta_latest(kv, node->key, node->keylen, +1); + priskv_resp_status resp = + priskv_pin_count_delta_latest(kv, node->key, node->keylen, +1, ttl_ms); if (resp == PRISKV_RESP_STATUS_OK) { kv->pin_stats.pin_ops++; + } else { + /* Count failed PIN attempts (e.g., NO_SUCH_KEY) */ + kv->pin_stats.pin_failed_ops++; } return resp; } @@ -934,6 +964,12 @@ uint64_t priskv_get_pin_ops(void *_kv) return kv ? kv->pin_stats.pin_ops : 0; } +uint64_t priskv_get_pin_failed_ops(void *_kv) +{ + priskv_kv *kv = (priskv_kv *)_kv; + return kv ? kv->pin_stats.pin_failed_ops : 0; +} + uint64_t priskv_get_unpin_ops(void *_kv) { priskv_kv *kv = (priskv_kv *)_kv; @@ -946,6 +982,8 @@ uint64_t priskv_get_unpin_not_closed(void *_kv) return kv ? kv->pin_stats.unpin_not_closed : 0; } +/* removed: priskv_get_pinop_invariant_violations; unified into unpin_not_closed */ + int priskv_value_addr_offset(void *_kv, uint8_t *val, uint64_t *addr_offset) { priskv_kv *kv = _kv; @@ -1165,8 +1203,22 @@ void priskv_clear_expired_kv(int fd, void *opaque, uint32_t events) priskv_hash_head *hash_head = &kv->hash_heads[i]; pthread_spin_lock(&hash_head->lock); + /* Single pass per bucket: collect expired keys and clean up expired pin TTLs */ list_for_each_safe (&hash_head->head, keynode, tmp, entry) { + /* TTL cleanup for non-expired keys: node level TTL expired set pin_count to 0 */ + pthread_spin_lock(&keynode->lock); + if (keynode->pin_count > 0) { + int64_t elapsed_ms = priskv_time_elapsed_ms(keynode->pin_ttl, now); + if (elapsed_ms > 0) { + /* TTL expired: clear pin_count and reset pin_ttl to invalid */ + keynode->pin_count = 0; + keynode->pin_ttl.tv_sec = -1; + keynode->pin_ttl.tv_usec = -1; + } + } + pthread_spin_unlock(&keynode->lock); if (priskv_key_timeout(keynode, now)) { + /* Collect expired KV for cleanup outside of the bucket lock */ list_del(&keynode->entry); list_add_tail(&expired_kv, &keynode->entry); kv->expire_routine_statics.expire_kv_count++; diff --git a/server/kv.h b/server/kv.h index a6544c8..061e654 100644 --- a/server/kv.h +++ b/server/kv.h @@ -42,6 +42,8 @@ struct priskv_transport_conn; typedef struct priskv_transport_conn priskv_transport_conn; #define PRISKV_KV_DEFAULT_EXPIRE_ROUTINE_INTERVAL 600 +/* Default TTL (ms) for each pin operation if not otherwise specified */ +#define PRISKV_DEFAULT_PIN_TTL_MS 60000 void *priskv_new_kv(uint8_t *key_base, uint8_t *value_base, int shm_fd, uint64_t shm_len, uint32_t max_keys, uint16_t max_key_length, uint32_t value_block_size, @@ -159,33 +161,20 @@ void priskv_key_serialize_exit(struct priskv_tiering_req *completed_req); int priskv_alloc_node_private(void *_kv, uint8_t *key, uint16_t keylen, uint8_t **val, uint32_t alloc_length, uint64_t timeout, void **_keynode); int priskv_publish_node(void *_kv, void *_keynode); -/* Atomically publish and optionally pin on publish (for SEAL with PIN). */ -int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish); +/* Atomically publish and optionally pin; ttl_ms == 0 uses the default TTL. */ +int priskv_publish_node_with_pin(void *_kv, void *_keynode, bool pin_on_publish, uint64_t ttl_ms); int priskv_drop_node(void *_kv, void *_keynode); int priskv_key_unpin_latest(void *_kv, void *_keynode); -/* Pin on the latest version of the key corresponding to keynode */ -int priskv_key_pin_latest(void *_kv, void *_keynode); +/* Pin on the latest version; ttl_ms==0 uses the default TTL */ +int priskv_key_pin_latest(void *_kv, void *_keynode, uint64_t ttl_ms); /* Pin/Unpin observability */ uint64_t priskv_get_pin_ops(void *_kv); +uint64_t priskv_get_pin_failed_ops(void *_kv); uint64_t priskv_get_unpin_ops(void *_kv); uint64_t priskv_get_unpin_not_closed(void *_kv); -/* - * TODO(wangyi): PinManager APIs & metrics - * - Define PinManager lifecycle APIs: - * - priskv_pin_manager_init/_destroy - * - priskv_pin_register(key, keylen, ttl_ms, origin, request_id) - * - priskv_pin_remove(key, keylen) - * - priskv_pin_ttl_cleanup_tick() scheduled via timerfd - * - Define metrics getters for info panel: - * - priskv_get_pin_ttl_active() - * - priskv_get_pin_ttl_expired() - * - priskv_get_pin_ttl_cleanup_ops() - * - priskv_get_pin_ttl_orphaned() - */ - #if defined(__cplusplus) } #endif diff --git a/server/memory.h b/server/memory.h index 6581738..a86fdf3 100644 --- a/server/memory.h +++ b/server/memory.h @@ -47,6 +47,7 @@ typedef struct priskv_key { pthread_spinlock_t lock; uint32_t refcnt; uint32_t pin_count; /* number of active pins; protect from eviction */ + struct timeval pin_ttl; /* node-level pin TTL; valid when pin_count > 0 */ /* TODO(wangyi): Guard against pin_count overflow * - Add warnings when pin_count approaches UINT32_MAX to prevent silent wrap-around. * - Consider switching to 64-bit counters if workload can accumulate many long-lived pins. diff --git a/server/test/Makefile b/server/test/Makefile index 757f051..ce8a819 100644 --- a/server/test/Makefile +++ b/server/test/Makefile @@ -11,6 +11,7 @@ TEST_MEMORY = test-memory TEST_ACL = test-acl TEST_KV_EXPIRE_ROUTINE = test-kv-expire-routine TEST_BE_REDIS = test-be-redis +TEST_KV_PIN_TTL = test-kv-pin-ttl CFLAGS = -fPIC -Wall -g -O0 -I .. -I ../../include -D_GNU_SOURCE -Wshadow -Wformat=2 -Wwrite-strings -fstack-protector-strong -Wnull-dereference -Wunreachable-code -lpthread -ldl FMT = clang-format-19 @@ -35,7 +36,7 @@ endif .PHONY: $(TEST_BUDDY) ${TEST_BUDDY_MT} $(TEST_SLAB) $(TEST_SLAB_MT) $(TEST_KV) $(TEST_KV_MT) $(TEST_TRANSPORT) $(TEST_MEMORY) $(TEST_ACL) $(TEST_KV_EXPIRE_ROUTINE) $(TEST_BE_REDIS) OBJS = ../memory.o ../kv.o ../slab.o ../crc.o ../acl.o -all: $(TEST_BUDDY) ${TEST_BUDDY_MT} $(TEST_SLAB) $(TEST_SLAB_MT) $(TEST_KV) $(TEST_KV_MT) $(TEST_TRANSPORT) $(TEST_MEMORY) $(TEST_ACL) $(TEST_KV_EXPIRE_ROUTINE) $(TEST_BE_REDIS) +all: $(TEST_BUDDY) ${TEST_BUDDY_MT} $(TEST_SLAB) $(TEST_SLAB_MT) $(TEST_KV) $(TEST_KV_MT) $(TEST_TRANSPORT) $(TEST_MEMORY) $(TEST_ACL) $(TEST_KV_EXPIRE_ROUTINE) $(TEST_BE_REDIS) $(TEST_KV_PIN_TTL) $(TEST_BUDDY): $(OBJS) $(CC) test_buddy.c ../buddy.c $(CFLAGS) -o $(TEST_BUDDY) @@ -87,6 +88,15 @@ $(TEST_KV_EXPIRE_ROUTINE): @echo "UCX not available, skip $(TEST_KV_EXPIRE_ROUTINE)" endif +ifeq ($(shell pkg-config --exists ucx; echo $$?),0) +$(TEST_KV_PIN_TTL): $(OBJS) + $(CC) test_kv_pin_ttl.c ../../lib/workqueue.c ../../lib/threads.c ../../lib/event.c ../../lib/ucx.c ../../lib/config.c ../memory.c ../kv.c ../slab.c ../buddy.c ../crc.c ../../lib/log.c ../backend/backend.c ../transport/transport.c ../transport/rdma.c ../transport/ucx.c ../tiering.c ../acl.c $(CFLAGS) $(UCX_CFLAGS) -o $(TEST_KV_PIN_TTL) -lmount -lpthread -lrdmacm -libverbs $(UCX_LIBS) +else +$(TEST_KV_PIN_TTL): + @echo "UCX not available, skip $(TEST_KV_PIN_TTL)" +endif + + $(TEST_BE_REDIS): $(CC) test_be_redis.c ../../lib/log.c ../../lib/event.c ../../lib/workqueue.c ../../lib/threads.c ../backend/backend.c ../backend/be_redis.c $(CFLAGS) -o $(TEST_BE_REDIS) -levent -lhiredis diff --git a/server/test/test_kv_pin_ttl.c b/server/test/test_kv_pin_ttl.c new file mode 100644 index 0000000..2f1dd53 --- /dev/null +++ b/server/test/test_kv_pin_ttl.c @@ -0,0 +1,316 @@ +// Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +// Licensed under the Apache License, Version 2.0 + +#include +#include +#include +#include +#include + +#include "kv.h" +#include "priskv-threads.h" +#include "priskv-protocol.h" +#include "memory.h" +#include "buddy.h" + +/* Create a small KV for tests; caller frees key_base/value_base and destroys kv */ +static void *create_kv(uint32_t max_keys, uint16_t max_key_length, uint32_t value_block_size, + uint64_t value_blocks, uint8_t **key_base_out, uint8_t **value_base_out) +{ + uint8_t *kbase = calloc(max_keys, priskv_mem_key_size(max_key_length)); + uint8_t *vbase = calloc(1, priskv_buddy_mem_size(value_blocks, value_block_size)); + void *kv = priskv_new_kv(kbase, vbase, -1, 0, max_keys, max_key_length, value_block_size, + value_blocks, NULL /* mf_ctx */); + if (!kv) { + free(kbase); + free(vbase); + *key_base_out = NULL; + *value_base_out = NULL; + return NULL; + } + *key_base_out = kbase; + *value_base_out = vbase; + return kv; +} + +/* Run TTL scenarios (immediate unpin, TTL expiry, default TTL) on a fresh KV */ +static int test_ttl_suite(void) +{ + /* thread pool for expire routine */ + uint8_t iothreads = 1, bgthreads = 1; + priskv_threadpool *tp = priskv_threadpool_create("test", iothreads, bgthreads, 0); + assert(tp); + priskv_thread *bgthread = priskv_threadpool_get_bgthread(tp, 0); + assert(bgthread); + + /* KV capacity */ + uint8_t *key_base = NULL, *value_base = NULL, *val_ptr = NULL; + void *kv = create_kv(1024, 128, 4096, 1024 * 4ULL, &key_base, &value_base); + if (!kv) { + printf("TEST KV PIN TTL: create kv [FAILED]\n"); + return 1; + } + + /* run expire routine in bg thread every 1s for TTL cleanup */ + priskv_set_expire_routine_interval(kv, 1 /* seconds */); + priskv_expire_routine(bgthread, kv); + + const uint8_t *key = (const uint8_t *)"kpin"; + const uint16_t keylen = 4; + const uint32_t alloc_len = 64; + void *keynode = NULL; + int ret; + + /* Scenario A: publish with short TTL and immediately unpin succeeds */ + ret = priskv_alloc_node_private(kv, (uint8_t *)key, keylen, &val_ptr, alloc_len, + PRISKV_KEY_MAX_TIMEOUT, &keynode); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: alloc private node [FAILED], ret=%d\n", ret); + goto fail; + } + ret = priskv_publish_node_with_pin(kv, keynode, true /* pin_on_publish */, 300 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: publish with pin [FAILED], ret=%d\n", ret); + goto fail; + } + ret = priskv_key_unpin_latest(kv, keynode); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: unpin after publish [FAILED], ret=%d\n", ret); + goto fail; + } + if (priskv_get_unpin_not_closed(kv) != 0) { + printf("TEST KV PIN TTL: unpin_not_closed changed unexpectedly [FAILED]\n"); + goto fail; + } + + /* Scenario B: TTL expiry cleans pinops, later UNPIN reports NOT_CLOSED */ + ret = priskv_alloc_node_private(kv, (uint8_t *)key, keylen, &val_ptr, alloc_len, + PRISKV_KEY_MAX_TIMEOUT, &keynode); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: alloc private node (B) [FAILED], ret=%d\n", ret); + goto fail; + } + ret = priskv_publish_node_with_pin(kv, keynode, true /* pin_on_publish */, 300 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: publish with pin (B) [FAILED], ret=%d\n", ret); + goto fail; + } + /* Wait for TTL (300ms) + expire tick (1s) */ + usleep(1600 * 1000); + ret = priskv_key_unpin_latest(kv, keynode); + if (ret != PRISKV_RESP_STATUS_UNPIN_NOT_CLOSED) { + printf("TEST KV PIN TTL: unpin after TTL expiry should be NOT_CLOSED [FAILED], ret=%d\n", + ret); + goto fail; + } + + /* Scenario C: ttl_ms == 0 uses default TTL path; just sanity-pin then unpin */ + ret = priskv_key_pin_latest(kv, keynode, 0 /* default TTL */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: pin with default TTL [FAILED], ret=%d\n", ret); + goto fail; + } + ret = priskv_key_unpin_latest(kv, keynode); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: unpin after default TTL pin [FAILED], ret=%d\n", ret); + goto fail; + } + + /* Scenario D: re-pin with shorter TTL should NOT shorten pin_ttl (use max) */ + ret = priskv_alloc_node_private(kv, (uint8_t *)key, keylen, &val_ptr, alloc_len, + PRISKV_KEY_MAX_TIMEOUT, &keynode); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: alloc private node (D) [FAILED], ret=%d\n", ret); + goto fail; + } + /* Publish with long TTL = 5000ms */ + ret = priskv_publish_node_with_pin(kv, keynode, true /* pin_on_publish */, 5000 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: publish with long pin (D) [FAILED], ret=%d\n", ret); + goto fail; + } + /* Re-pin with short TTL = 300ms; pin_ttl should remain at max(5000ms, now+300ms) */ + ret = priskv_key_pin_latest(kv, keynode, 300 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: re-pin with short TTL (D) [FAILED], ret=%d\n", ret); + goto fail; + } + /* Wait for expire tick (1s) so that if pin_ttl were shortened to 300ms, cleanup would reset */ + usleep(1600 * 1000); + /* Expect UNPIN still succeeds because original longer TTL keeps pin active */ + ret = priskv_key_unpin_latest(kv, keynode); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN TTL: unpin after re-pin short (should be OK) [FAILED], ret=%d\n", ret); + goto fail; + } + + priskv_destroy_kv(kv); + free(key_base); + free(value_base); + printf("TEST KV PIN TTL: [OK]\n"); + return 0; + +fail: + priskv_destroy_kv(kv); + free(key_base); + free(value_base); + return 1; +} + +/* Run pin inherit semantics on a fresh KV */ +static int test_inherit_suite(void) +{ + uint8_t *key_base = NULL, *value_base = NULL, *val_ptr = NULL; + void *kv = create_kv(256, 64, 4096, 256 * 2ULL, &key_base, &value_base); + if (!kv) { + printf("TEST KV PIN INHERIT: create kv [FAILED]\n"); + return 1; + } + + int ret; + void *kn1 = NULL, *kn2 = NULL; + const uint8_t *key = (const uint8_t *)"kver"; + const uint16_t keylen = 4; + const uint32_t alloc_len = 32; + + /* v1: ALLOC + PUBLISH(pin) */ + ret = priskv_alloc_node_private(kv, (uint8_t *)key, keylen, &val_ptr, alloc_len, + PRISKV_KEY_MAX_TIMEOUT, &kn1); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN INHERIT: alloc v1 [FAILED], ret=%d\n", ret); + goto fail; + } + ret = priskv_publish_node_with_pin(kv, kn1, true /* pin_on_publish */, 10000 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN INHERIT: publish v1 with pin [FAILED], ret=%d\n", ret); + goto fail; + } + + /* v2: ALLOC + PUBLISH(pin) */ + ret = priskv_alloc_node_private(kv, (uint8_t *)key, keylen, &val_ptr, alloc_len, + PRISKV_KEY_MAX_TIMEOUT, &kn2); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN INHERIT: alloc v2 [FAILED], ret=%d\n", ret); + goto fail; + } + ret = priskv_publish_node_with_pin(kv, kn2, true /* pin_on_publish */, 10000 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN INHERIT: publish v2 with pin [FAILED], ret=%d\n", ret); + goto fail; + } + + /* UNPIN #1 */ + ret = priskv_key_unpin_latest(kv, kn2); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN INHERIT: first unpin [FAILED], ret=%d\n", ret); + goto fail; + } + + /* UNPIN #2 */ + ret = priskv_key_unpin_latest(kv, kn2); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PIN INHERIT: second unpin [FAILED], ret=%d\n", ret); + goto fail; + } + + /* Extra UNPIN should report NOT_CLOSED */ + ret = priskv_key_unpin_latest(kv, kn2); + if (ret != PRISKV_RESP_STATUS_UNPIN_NOT_CLOSED) { + printf("TEST KV PIN INHERIT: extra unpin should be NOT_CLOSED [FAILED], ret=%d\n", ret); + goto fail; + } + + priskv_destroy_kv(kv); + free(key_base); + free(value_base); + printf("TEST KV PIN INHERIT: [OK]\n"); + return 0; + +fail: + priskv_destroy_kv(kv); + free(key_base); + free(value_base); + return 1; +} + +/* Run pinned-eviction semantics on a fresh KV */ +static int test_pinned_eviction_suite(void) +{ + /* Tiny KV to force pressure */ + uint8_t *key_base = NULL, *value_base = NULL, *val_ptr = NULL; + void *kv = create_kv(8 /* max_keys */, 32 /* max_key_len */, 1024 /* vblk_size */, + 8 /* vblks */, &key_base, &value_base); + if (!kv) { + printf("TEST KV PINNED EVICTION: create kv [FAILED]\n"); + return 1; + } + + int ret; + void *kn_pinned = NULL, *kn_tmp = NULL, *kn_check = NULL; + const uint8_t *kpin = (const uint8_t *)"kpin"; + const uint8_t *ktemp = (const uint8_t *)"ktemp"; + const uint8_t *kbig = (const uint8_t *)"kbig"; + uint8_t *val_check = NULL; + uint32_t vlen_check = 0; + + /* Insert pinned key with a small value, then pin it */ + ret = priskv_set_key(kv, (uint8_t *)kpin, 4, &val_ptr, 512, PRISKV_KEY_MAX_TIMEOUT, &kn_pinned); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PINNED EVICTION: set pinned key [FAILED], ret=%d\n", ret); + goto fail; + } + priskv_set_key_end(kn_pinned); + ret = priskv_key_pin_latest(kv, kn_pinned, 10000 /* ms */); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PINNED EVICTION: pin pinned key [FAILED], ret=%d\n", ret); + goto fail; + } + + /* Insert a temporary unpinned key to be evicted */ + ret = priskv_set_key(kv, (uint8_t *)ktemp, 5, &val_ptr, 1024, PRISKV_KEY_MAX_TIMEOUT, &kn_tmp); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PINNED EVICTION: set temp key [FAILED], ret=%d\n", ret); + goto fail; + } + priskv_set_key_end(kn_tmp); + + /* Now try to insert a big value that forces eviction */ + (void)priskv_set_key(kv, (uint8_t *)kbig, 4, &val_ptr, 1024 * 6, PRISKV_KEY_MAX_TIMEOUT, + &kn_tmp); + + /* Verify pinned key still exists */ + ret = priskv_get_key(kv, (uint8_t *)kpin, 4, &val_check, &vlen_check, &kn_check); + if (ret != PRISKV_RESP_STATUS_OK) { + printf("TEST KV PINNED EVICTION: pinned key missing after pressure [FAILED], ret=%d\n", + ret); + goto fail; + } + priskv_get_key_end(kn_check); + + /* Cleanup */ + (void)priskv_key_unpin_latest(kv, kn_pinned); + priskv_destroy_kv(kv); + free(key_base); + free(value_base); + printf("TEST KV PINNED EVICTION: [OK]\n"); + return 0; + +fail: + (void)priskv_key_unpin_latest(kv, kn_pinned); + priskv_destroy_kv(kv); + free(key_base); + free(value_base); + return 1; +} + +int main() +{ + if (test_ttl_suite() != 0) + return 1; + if (test_inherit_suite() != 0) + return 1; + if (test_pinned_eviction_suite() != 0) + return 1; + printf("TEST KV PIN TTL + INHERIT: [OK]\n"); + return 0; +} diff --git a/server/test/test_transport.c b/server/test/test_transport.c index d41e41e..3761f93 100644 --- a/server/test/test_transport.c +++ b/server/test/test_transport.c @@ -470,15 +470,18 @@ typedef struct seal_req_arg { static void *seal_with_pin_thread(void *arg) { seal_req_arg *a = (seal_req_arg *)arg; - uint8_t req_buf[128]; + const size_t req_size = sizeof(priskv_request) + sizeof(uint64_t); + uint8_t req_buf[req_size]; priskv_request *req = (priskv_request *)req_buf; - memset(req, 0, sizeof(req_buf)); + memset(req, 0, req_size); req->command = htobe16(PRISKV_COMMAND_SEAL); req->flags = htobe32(PRISKV_REQ_FLAG_PIN_ON_SEAL); + req->nsgl = htobe16(0); + req->key_length = htobe16(sizeof(uint64_t)); uint64_t be_token = htobe64(a->token_host); memcpy((uint8_t *)req + sizeof(priskv_request), &be_token, sizeof(uint64_t)); pthread_barrier_wait(a->barrier); - priskv_transport_handle_recv(a->conn, req, sizeof(priskv_request) + sizeof(uint64_t)); + priskv_transport_handle_recv(a->conn, req, (uint16_t)req_size); return NULL; } @@ -557,15 +560,18 @@ static void test_kv_transport_concurrent_seal_pin(void *kv) static void *seal_with_pin_thread_stress(void *arg) { seal_req_arg *a = (seal_req_arg *)arg; - uint8_t req_buf[128]; + const size_t req_size = sizeof(priskv_request) + sizeof(uint64_t); + uint8_t req_buf[req_size]; priskv_request *req = (priskv_request *)req_buf; - memset(req, 0, sizeof(req_buf)); + memset(req, 0, req_size); req->command = htobe16(PRISKV_COMMAND_SEAL); req->flags = htobe32(PRISKV_REQ_FLAG_PIN_ON_SEAL); + req->nsgl = htobe16(0); + req->key_length = htobe16(sizeof(uint64_t)); uint64_t be_token = htobe64(a->token_host); memcpy((uint8_t *)req + sizeof(priskv_request), &be_token, sizeof(uint64_t)); pthread_barrier_wait(a->barrier); - priskv_transport_handle_recv(a->conn, req, sizeof(priskv_request) + sizeof(uint64_t)); + priskv_transport_handle_recv(a->conn, req, (uint16_t)req_size); return NULL; } diff --git a/server/transport/transport.c b/server/transport/transport.c index 667d9ed..8edb376 100644 --- a/server/transport/transport.c +++ b/server/transport/transport.c @@ -287,6 +287,7 @@ int priskv_transport_handle_recv(priskv_transport_conn *conn, priskv_request *re uint32_t flags = be32toh(req->flags); uint16_t nsgl = be16toh(req->nsgl); uint64_t timeout = be64toh(req->timeout); + uint64_t pin_ttl_ms = be64toh(req->pin_ttl_ms); uint32_t alloc_length = be32toh(req->alloc_length); uint8_t *key; uint16_t keylen; @@ -592,14 +593,14 @@ int priskv_transport_handle_recv(priskv_transport_conn *conn, priskv_request *re PRISKV_RESP_STATUS_PERMISSION_DENIED, 0, 0, 0); break; } - /* Atomically publish and optionally pin within KV to remove the window. */ - status = priskv_publish_node_with_pin(conn->kv, keynode, - (flags & PRISKV_REQ_FLAG_PIN_ON_SEAL) != 0); - /* TODO(wangyi): PinTTL register on SEAL - * - If protocol provides per-request TTL (e.g., pin_ttl_ms), register a - * PinOperator with PinManager for this key to ensure eventual cleanup. - * - Fallback to server default TTL when not provided. - */ + /* Atomically publish and optionally pin; treat req.timeout as the TTL (ms) for this + * pin. */ + /* Use pin_ttl_ms; 0 means default TTL. Only read TTL when pin_on_seal is set. */ + status = priskv_publish_node_with_pin( + conn->kv, keynode, (flags & PRISKV_REQ_FLAG_PIN_ON_SEAL) != 0, + (flags & PRISKV_REQ_FLAG_PIN_ON_SEAL) ? pin_ttl_ms : 0); + /* TTL registration is handled in the KV layer (publish critical section); no need to + * duplicate here. */ priskv_transport_token_del(conn, token); ret = driver->send_response(conn, req->request_id, status, 0, 0, 0); } @@ -619,7 +620,8 @@ int priskv_transport_handle_recv(priskv_transport_conn *conn, priskv_request *re * the reference acquired by priskv_get_key. */ if (flags & PRISKV_REQ_FLAG_PIN_ON_ACQUIRE) { - priskv_resp_status presp = priskv_key_pin_latest(conn->kv, keynode); + /* Use pin_ttl_ms; 0 means default TTL. */ + priskv_resp_status presp = priskv_key_pin_latest(conn->kv, keynode, pin_ttl_ms); if (presp != PRISKV_RESP_STATUS_OK) { /* Atomicity: no token created, drop our reference and return pin's status */ priskv_get_key_end(keynode); @@ -640,10 +642,7 @@ int priskv_transport_handle_recv(priskv_transport_conn *conn, priskv_request *re 0, 0, 0); break; } - /* TODO(wangyi): PinTTL register on ACQUIRE - * - Register PinOperator with PinManager using request-scoped or default TTL. - * - Associate optional request_id for observability. - */ + /* TTL registration is handled in the KV layer (pin path); no need to duplicate here. */ status = priskv_value_addr_offset(conn->kv, val, &addr_offset); ret = driver->send_response(conn, req->request_id, status, valuelen, addr_offset, token); @@ -676,26 +675,22 @@ int priskv_transport_handle_recv(priskv_transport_conn *conn, priskv_request *re PRISKV_RESP_STATUS_PERMISSION_DENIED, 0, 0, 0); break; } - /* Enforce RELEASE+UNPIN all-or-nothing semantics when requested: - * - If UNPIN_ON_RELEASE is set and unpin fails (e.g., NO_SUCH_KEY / UNPIN_NOT_CLOSED), - * do NOT release the ACQUIRE reference or delete the token. Client can retry. - * - Otherwise, proceed to release and delete token, returning OK. + /* RELEASE semantics when UNPIN is requested: + * - If PRISKV_REQ_FLAG_UNPIN_ON_RELEASE is set and unpin fails (e.g., NO_SUCH_KEY / + * UNPIN_NOT_CLOSED), we currently release the ACQUIRE reference and delete the + * token as a best-effort cleanup, and return the unpin status to the client. + * - Otherwise, we release the reference and delete the token, returning OK. + * Note: This favors simplicity over retry semantics; adjust if caller requires a + * strict all-or-nothing behavior. */ priskv_resp_status resp = PRISKV_RESP_STATUS_OK; if (flags & PRISKV_REQ_FLAG_UNPIN_ON_RELEASE) { resp = priskv_key_unpin_latest(conn->kv, keynode); - if (resp != PRISKV_RESP_STATUS_OK) { - /* unpin failed but still cleanup token and reference */ - priskv_get_key_end(keynode); - priskv_transport_token_del(conn, token); - ret = driver->send_response(conn, req->request_id, resp, 0, 0, 0); - break; - } } /* Either UNPIN succeeded or not requested: finish RELEASE */ priskv_get_key_end(keynode); priskv_transport_token_del(conn, token); - ret = driver->send_response(conn, req->request_id, PRISKV_RESP_STATUS_OK, 0, 0, 0); + ret = driver->send_response(conn, req->request_id, resp, 0, 0, 0); } break; case PRISKV_COMMAND_DROP: diff --git a/server/transport/ucx.c b/server/transport/ucx.c index 7a6c910..e9bed1b 100644 --- a/server/transport/ucx.c +++ b/server/transport/ucx.c @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -590,7 +591,7 @@ static int priskv_ucx_listen_one(char *addr, int port, void *kv, priskv_transpor goto out_free_res; } - ret = listen(listenfd, 0); + ret = listen(listenfd, SOMAXCONN); if (ret < 0) { priskv_log_error("UCX: listen failed: %m\n"); ret = -1;