feat: add SCANPREFIX/KEYSPREFIX/SCANRANGE/KEYSRANGE

kvs_protocol_resp.c

@@ -275,6 +275,21 @@ int resp_build_value(const resp_value_t *v, uint8_t *out, size_t cap) {
         if (write_crlf(&p, end) < 0) return -1;
         break;
 
+    case RESP_T_ARRAY:
+        if (write_bytes(&p, end, "*", 1) < 0) return -1;
+        if (write_u32_ascii(&p, end, v->array.count) < 0) return -1;
+        if (write_crlf(&p, end) < 0) return -1;
+        for (uint32_t i = 0; i < v->array.count; ++i) {
+            if (write_bytes(&p, end, "$", 1) < 0) return -1;
+            if (write_u32_ascii(&p, end, v->array.items[i].len) < 0) return -1;
+            if (write_crlf(&p, end) < 0) return -1;
+            if (v->array.items[i].len > 0 && v->array.items[i].ptr) {
+                if (write_bytes(&p, end, v->array.items[i].ptr, v->array.items[i].len) < 0) return -1;
+            }
+            if (write_crlf(&p, end) < 0) return -1;
+        }
+        break;
+
     default:
         return -1;
     }
@@ -285,8 +300,8 @@ int resp_build_value(const resp_value_t *v, uint8_t *out, size_t cap) {
 /* helpers */
 resp_value_t resp_simple(const char *s) {
     resp_value_t v;
+    memset(&v, 0, sizeof(v));
     v.type = RESP_T_SIMPLE_STR;
-    v.i64 = 0;
     v.bulk.ptr = (const uint8_t*)s;
     v.bulk.len = (uint32_t)strlen(s);
     return v;
@@ -294,8 +309,8 @@ resp_value_t resp_simple(const char *s) {
 
 resp_value_t resp_error(const char *s) {
     resp_value_t v;
+    memset(&v, 0, sizeof(v));
     v.type = RESP_T_ERROR;
-    v.i64 = 0;
     v.bulk.ptr = (const uint8_t*)s;
     v.bulk.len = (uint32_t)strlen(s);
     return v;
@@ -303,17 +318,16 @@ resp_value_t resp_error(const char *s) {
 
 resp_value_t resp_int(int64_t x) {
     resp_value_t v;
+    memset(&v, 0, sizeof(v));
     v.type = RESP_T_INTEGER;
     v.i64 = x;
-    v.bulk.ptr = NULL;
-    v.bulk.len = 0;
     return v;
 }
 
 resp_value_t resp_bulk(const uint8_t *p, uint32_t n) {
     resp_value_t v;
+    memset(&v, 0, sizeof(v));
     v.type = RESP_T_BULK_STR;
-    v.i64 = 0;
     v.bulk.ptr = p;
     v.bulk.len = n;
     return v;
@@ -321,13 +335,31 @@ resp_value_t resp_bulk(const uint8_t *p, uint32_t n) {
 
 resp_value_t resp_nil(void) {
     resp_value_t v;
+    memset(&v, 0, sizeof(v));
     v.type = RESP_T_NIL;
-    v.i64 = 0;
-    v.bulk.ptr = NULL;
-    v.bulk.len = 0;
     return v;
 }
 
+resp_value_t resp_array(resp_slice_t *items, uint32_t count, int owns_items) {
+    resp_value_t v;
+    memset(&v, 0, sizeof(v));
+    v.type = RESP_T_ARRAY;
+    v.array.items = items;
+    v.array.count = count;
+    v.array.owns_items = owns_items ? 1 : 0;
+    return v;
+}
+
+void resp_value_release(resp_value_t *v) {
+    if (!v) return;
+    if (v->type == RESP_T_ARRAY && v->array.owns_items && v->array.items) {
+        kvs_free(v->array.items);
+    }
+    memset(v, 0, sizeof(*v));
+    v->type = RESP_T_NIL;
+    return;
+}
+
 /* ----------------- dispatcher (minimal) ----------------- */
 
 static int expect_argv(const resp_cmd_t *cmd, uint32_t n) {
@@ -337,7 +369,7 @@ static int expect_argv(const resp_cmd_t *cmd, uint32_t n) {
 
 const char *command[] = {
 	"PING", "AUTH",
-	"SET", "GET", "DEL", "SETEX", "MOD", "EXIST",
+	"SET", "GET", "DEL", "SETEX", "MOD", "EXIST", "SCANPREFIX", "KEYSPREFIX", "SCANRANGE", "KEYSRANGE",
 	"RSET", "RGET", "RDEL", "RMOD", "REXIST",
 	"HSET", "HGET", "HDEL", "HMOD", "HEXIST",
 	"SAVE", "SSYNC", "SREADY", "MEMPRINT"
@@ -382,6 +414,120 @@ static int match_auth_password(const resp_slice_t *password) {
     return memcmp(global_cfg.redis_auth_password, password->ptr, expect_len) == 0;
 }
 
+static int prefix_match(const uint8_t *key, uint32_t key_len,
+                        const uint8_t *prefix, uint32_t prefix_len) {
+    if (!key || !prefix) return 0;
+    if (key_len < prefix_len) return 0;
+    return memcmp(key, prefix, prefix_len) == 0;
+}
+
+static rbtree_node *rbtree_lower_bound(rbtree *T, const uint8_t *key, uint32_t keylen) {
+    rbtree_node *cur;
+    rbtree_node *best;
+
+    if (!T || !key) return NULL;
+
+    cur = T->root;
+    best = T->nil;
+    while (cur != T->nil) {
+        const uint8_t *cur_key = kvs_rbtree_node_key_ptr(cur);
+        int cmp = kvs_keycmp(cur_key, cur->key_len, key, keylen);
+        if (cmp < 0) {
+            cur = cur->right;
+        } else {
+            best = cur;
+            cur = cur->left;
+        }
+    }
+
+    return best;
+}
+
+static int append_scan_item(resp_slice_t **items, uint32_t *count, uint32_t *cap,
+                            const uint8_t *ptr, uint32_t len) {
+    resp_slice_t *new_items;
+    uint32_t new_cap;
+
+    if (!items || !count || !cap) return -1;
+
+    if (*count >= *cap) {
+        new_cap = (*cap == 0) ? 16u : (*cap * 2u);
+        new_items = (resp_slice_t *)kvs_malloc(sizeof(resp_slice_t) * (size_t)new_cap);
+        if (!new_items) return -1;
+        if (*items && *count > 0) {
+            memcpy(new_items, *items, sizeof(resp_slice_t) * (size_t)(*count));
+            kvs_free(*items);
+        }
+        *items = new_items;
+        *cap = new_cap;
+    }
+
+    (*items)[*count].ptr = ptr;
+    (*items)[*count].len = len;
+    (*count)++;
+    return 0;
+}
+
+static int build_rbtree_scan_response(const uint8_t *begin, uint32_t begin_len,
+                                      const uint8_t *end, uint32_t end_len,
+                                      int prefix_mode, int keys_only,
+                                      resp_value_t *out_value) {
+    resp_slice_t *items = NULL;
+    uint32_t count = 0;
+    uint32_t cap = 0;
+    rbtree_node *node;
+    uint64_t now_ms;
+
+    if (!out_value) return -1;
+
+    now_ms = kvs_now_ms();
+    if (begin && begin_len > 0) {
+        node = rbtree_lower_bound(&global_rbtree, begin, begin_len);
+    } else if (global_rbtree.root != global_rbtree.nil) {
+        node = rbtree_mini(&global_rbtree, global_rbtree.root);
+    } else {
+        node = global_rbtree.nil;
+    }
+
+    while (node && node != global_rbtree.nil) {
+        const uint8_t *key = kvs_rbtree_node_key_ptr(node);
+        const uint8_t *value = kvs_rbtree_node_value_ptr(node);
+        rbtree_node *next = rbtree_successor(&global_rbtree, node);
+
+        if (prefix_mode) {
+            if (begin && begin_len > 0 && !prefix_match(key, node->key_len, begin, begin_len)) {
+                break;
+            }
+        } else if (end && end_len > 0) {
+            int end_cmp = kvs_keycmp(key, node->key_len, end, end_len);
+            if (end_cmp > 0) {
+                break;
+            }
+        }
+
+        if (node->expire_at_ms != 0 && node->expire_at_ms <= now_ms) {
+            (void)ttl_delete_schedule(key, node->key_len);
+            node = next;
+            continue;
+        }
+
+        if (append_scan_item(&items, &count, &cap, key, node->key_len) < 0) {
+            if (items) kvs_free(items);
+            return -1;
+        }
+        if (!keys_only &&
+            append_scan_item(&items, &count, &cap, value, node->value_len) < 0) {
+            if (items) kvs_free(items);
+            return -1;
+        }
+
+        node = next;
+    }
+
+    *out_value = resp_array(items, count, 1);
+    return 0;
+}
+
 
 /**
  * 输入：cmd
@@ -462,6 +608,44 @@ int resp_dispatch(const resp_cmd_t *cmd, resp_value_t *out_value) {
         return 0;
     }
 
+    case KVS_CMD_SCANPREFIX: {
+        if (cmd->argc != 2) { *out_value = resp_error("ERR wrong number of arguments for 'scanprefix'"); return -1; }
+        if (build_rbtree_scan_response(cmd->argv[1].ptr, cmd->argv[1].len,
+                                       NULL, 0, 1, 0, out_value) < 0) {
+            *out_value = resp_error("ERR internal error");
+        }
+        return 0;
+    }
+
+    case KVS_CMD_KEYSPREFIX: {
+        if (cmd->argc != 2) { *out_value = resp_error("ERR wrong number of arguments for 'keysprefix'"); return -1; }
+        if (build_rbtree_scan_response(cmd->argv[1].ptr, cmd->argv[1].len,
+                                       NULL, 0, 1, 1, out_value) < 0) {
+            *out_value = resp_error("ERR internal error");
+        }
+        return 0;
+    }
+
+    case KVS_CMD_SCANRANGE: {
+        if (cmd->argc != 3) { *out_value = resp_error("ERR wrong number of arguments for 'scanrange'"); return -1; }
+        if (build_rbtree_scan_response(cmd->argv[1].ptr, cmd->argv[1].len,
+                                       cmd->argv[2].ptr, cmd->argv[2].len,
+                                       0, 0, out_value) < 0) {
+            *out_value = resp_error("ERR internal error");
+        }
+        return 0;
+    }
+
+    case KVS_CMD_KEYSRANGE: {
+        if (cmd->argc != 3) { *out_value = resp_error("ERR wrong number of arguments for 'keysrange'"); return -1; }
+        if (build_rbtree_scan_response(cmd->argv[1].ptr, cmd->argv[1].len,
+                                       cmd->argv[2].ptr, cmd->argv[2].len,
+                                       0, 1, out_value) < 0) {
+            *out_value = resp_error("ERR internal error");
+        }
+        return 0;
+    }
+
     case KVS_CMD_DEL: {
         if (cmd->argc != 2) { *out_value = resp_error("ERR wrong number of arguments for 'del'"); return -1; }
         int r = kvs_rbtree_del(&global_rbtree, cmd->argv[1].ptr, cmd->argv[1].len);

kvs_protocol_resp.h

@@ -26,6 +26,10 @@ typedef enum {
 	KVS_CMD_SETEX,
 	KVS_CMD_MOD,
 	KVS_CMD_EXIST,
+	KVS_CMD_SCANPREFIX,
+	KVS_CMD_KEYSPREFIX,
+	KVS_CMD_SCANRANGE,
+	KVS_CMD_KEYSRANGE,
 	// rbtree
 	KVS_CMD_RSET,
 	KVS_CMD_RGET,
@@ -62,12 +66,20 @@ typedef enum resp_type_e {
     RESP_T_INTEGER,
     RESP_T_BULK_STR,
     RESP_T_NIL,          /* nil bulk string ($-1\r\n) */
+    RESP_T_ARRAY,
 } resp_type_t;
 
+typedef struct resp_array_s {
+    resp_slice_t *items;
+    uint32_t count;
+    int owns_items;
+} resp_array_t;
+
 typedef struct resp_value_s {
     resp_type_t type;
     int64_t i64;         /* for integer */
     resp_slice_t bulk;   /* for simple/error/bulk string (bytes) */
+    resp_array_t array;  /* for arrays */
 } resp_value_t;
 
 /* ----------------- parsing ----------------- */
@@ -86,6 +98,8 @@ resp_value_t resp_error(const char *s);  /* -s\r\n */
 resp_value_t resp_int(int64_t x);        /* :x\r\n */
 resp_value_t resp_bulk(const uint8_t *p, uint32_t n); /* $n\r\n...\r\n */
 resp_value_t resp_nil(void);             /* $-1\r\n */
+resp_value_t resp_array(resp_slice_t *items, uint32_t count, int owns_items);
+void resp_value_release(resp_value_t *v);
 
 /* ----------------- dispatcher ----------------- */
 /* Execute one command and return a RESP value.

kvstore.c

@@ -186,6 +186,51 @@ static int resp_value_encoded_len(const resp_value_t *v, size_t *out_len) {
         break;
     }
 
+    case RESP_T_ARRAY: {
+        char tmp[32];
+        int n;
+        size_t t = 0;
+
+        if (v->array.count > 0 && !v->array.items) {
+            return -1;
+        }
+
+        n = snprintf(tmp, sizeof(tmp), "%u", (unsigned)v->array.count);
+        if (n <= 0) {
+            return -1;
+        }
+
+        if (checked_size_add(1, (size_t)n, &t) < 0 ||
+            checked_size_add(t, 2, &t) < 0) {
+            return -1;
+        }
+
+        for (uint32_t i = 0; i < v->array.count; ++i) {
+            size_t item_len = 0;
+            int item_digits;
+
+            if (v->array.items[i].len > 0 && !v->array.items[i].ptr) {
+                return -1;
+            }
+
+            item_digits = snprintf(tmp, sizeof(tmp), "%u", (unsigned)v->array.items[i].len);
+            if (item_digits <= 0) {
+                return -1;
+            }
+
+            if (checked_size_add(1, (size_t)item_digits, &item_len) < 0 ||
+                checked_size_add(item_len, 2, &item_len) < 0 ||
+                checked_size_add(item_len, (size_t)v->array.items[i].len, &item_len) < 0 ||
+                checked_size_add(item_len, 2, &item_len) < 0 ||
+                checked_size_add(t, item_len, &t) < 0) {
+                return -1;
+            }
+        }
+
+        len = t;
+        break;
+    }
+
     default:
         return -1;
     }
@@ -328,11 +373,13 @@ int kvs_protocol(struct conn* conn){
                 uint8_t *resp_heap = NULL;
 
                 if (resp_value_encoded_len(&val, &resp_need) < 0) {
+                    resp_value_release(&val);
                     return -1;
                 }
 
                 resp_heap = (uint8_t *)kvs_malloc(resp_need);
                 if (!resp_heap) {
+                    resp_value_release(&val);
                     return -1;
                 }
 
@@ -340,6 +387,7 @@ int kvs_protocol(struct conn* conn){
                 if (resp_len < 0 ||
                     chain_buffer_append(&conn->wbuf, resp_heap, (size_t)resp_len) < 0) {
                     free(resp_heap);
+                    resp_value_release(&val);
                     return -1;
                 }
 
@@ -349,6 +397,7 @@ int kvs_protocol(struct conn* conn){
         }
 
         out_len += (size_t)resp_len;
+        resp_value_release(&val);
 
 		// __completed_cmd(request, consumed, 0);
 

kvstore.h

@@ -165,8 +165,11 @@ rbtree_node *rbtree_search(rbtree *T, const uint8_t *key, uint32_t keylen);
 typedef int (*kvs_rbtree_visit_cb)(rbtree_node *node, void *arg);
 rbtree_node *rbtree_search_with_visit(rbtree *T, const uint8_t *key, uint32_t keylen,
                                       kvs_rbtree_visit_cb cb, void *arg);
+rbtree_node *rbtree_mini(rbtree *T, rbtree_node *x);
+rbtree_node *rbtree_successor(rbtree *T, rbtree_node *x);
 const uint8_t *kvs_rbtree_node_key_ptr(const rbtree_node *node);
 const uint8_t *kvs_rbtree_node_value_ptr(const rbtree_node *node);
+int kvs_keycmp(const uint8_t *a, uint32_t alen, const uint8_t *b, uint32_t blen);
 
 int kvs_rbtree_save(iouring_ctx_t *uring, kvs_rbtree_t *inst, const char* filename);
 int kvs_rbtree_load(kvs_rbtree_t *inst, const char* filename);