redis scan command support

src/nc_core.h

@@ -63,6 +63,9 @@
 #define NC_EAGAIN   -2
 #define NC_ENOMEM   -3
 
+#define NC_MAX_NSERVER_BITS 12
+#define NC_MAX_NSERVER_MASK ((1<<NC_MAX_NSERVER_BITS) -1)
+
 /* reserved fds for std streams, log, stats fd, epoll etc. */
 #define RESERVED_FDS 32
 

src/nc_message.h

@@ -75,6 +75,7 @@ typedef enum msg_parse_result {
     ACTION( REQ_REDIS_PERSIST )                                                                     \
     ACTION( REQ_REDIS_PTTL )                                                                        \
     ACTION( REQ_REDIS_SORT )                                                                        \
+    ACTION( REQ_REDIS_SCAN )                                                                        \
     ACTION( REQ_REDIS_TOUCH )                                                                       \
     ACTION( REQ_REDIS_TTL )                                                                         \
     ACTION( REQ_REDIS_TYPE )                                                                        \
@@ -196,6 +197,7 @@ typedef enum msg_parse_result {
     ACTION( REQ_REDIS_GEOSEARCHSTORE)                                                               \
     ACTION( REQ_REDIS_EVAL )                   /* redis requests - eval */                          \
     ACTION( REQ_REDIS_EVALSHA )                                                                     \
+    ACTION( REQ_REDIS_SCRIPT)                                                                       \
     ACTION( REQ_REDIS_PING )                   /* redis requests - ping/quit */                     \
     ACTION( REQ_REDIS_QUIT)                                                                         \
     ACTION( REQ_REDIS_AUTH)                                                                         \
@@ -288,6 +290,7 @@ struct msg {
     uint32_t             nfrag_done;      /* # fragment done */
     uint64_t             frag_id;         /* id of fragmented message */
     struct msg           **frag_seq;      /* sequence of fragment message, map from keys to fragments*/
+    uint32_t             frag_multibulk_len; /* fragment response multibulk length */
 
     err_t                err;             /* errno on error? */
     unsigned             error:1;         /* error? */
@@ -300,6 +303,8 @@ struct msg {
     unsigned             fdone:1;         /* all fragments are done? */
     unsigned             swallow:1;       /* swallow response? */
     unsigned             redis:1;         /* redis? */
+
+    uint32_t             server_index;    /* the server index which the requstion should be forwarded */
 };
 
 TAILQ_HEAD(msg_tqh, msg);

src/nc_request.c

@@ -17,6 +17,7 @@
 
 #include <nc_core.h>
 #include <nc_server.h>
+#include <stdlib.h>
 
 struct msg *
 req_get(struct conn *conn)
@@ -471,6 +472,49 @@ req_make_reply(struct context *ctx, struct conn *conn, struct msg *req)
     return NC_OK;
 }
 
+/* If the user request is "scan 45066",
+ * the cursor 45066 in the request,
+ * we get server_index=45066 & NC_MAX_NSERVER_MASK=10,
+ * real_cursor = 45066>>NC_MAX_NSERVER_BITS = 11,
+ * and finally the request sent to the redis server will be "scan 00011".
+*/
+static void req_update_for_scan(struct server_pool *pool,struct msg *msg)
+{
+    if (msg->type != MSG_REQ_REDIS_SCAN) {
+        return;
+    }
+
+    ASSERT(array_n(msg->keys) > 0);
+
+    uint8_t *key;
+    uint32_t keylen;
+    struct keypos *kpos;
+    uint32_t idx;
+    unsigned long long cursor;
+    unsigned long real_cursor;
+    char arr[16];
+    char format[16];
+
+    kpos = array_get(msg->keys, 0);
+    key = kpos->start;
+    keylen = (uint32_t)(kpos->end - kpos->start);
+
+    if (keylen == 1 && key[0] == '0') {
+        idx=0;
+    }else{
+        cursor=strtoull((const char *)key,NULL,10);
+        idx = cursor & NC_MAX_NSERVER_MASK;
+        if (array_n(&pool->server)<=idx){
+            idx=0;
+        }
+        real_cursor = (cursor >> NC_MAX_NSERVER_BITS);
+        sprintf(format,"%%0%dd",keylen);
+        sprintf(arr,format,real_cursor);
+        nc_memcpy(key,arr,keylen);
+    }
+    msg->server_index=idx;
+}
+
 static bool
 req_filter(struct conn *conn, struct msg *msg)
 {
@@ -511,6 +555,8 @@ req_filter(struct conn *conn, struct msg *msg)
         msg->noforward = 1;
     }
 
+    req_update_for_scan(conn->owner,msg);
+
     return false;
 }
 
@@ -573,7 +619,7 @@ req_forward(struct context *ctx, struct conn *c_conn, struct msg *msg)
     key = kpos->start;
     keylen = (uint32_t)(kpos->end - kpos->start);
 
-    s_conn = server_pool_conn(ctx, c_conn->owner, key, keylen);
+    s_conn = server_pool_conn(ctx, c_conn->owner, msg, key, keylen);
     if (s_conn == NULL) {
         /*
          * Handle a failure to establish a new connection to a server,

src/nc_response.c

@@ -17,6 +17,7 @@
 
 #include <nc_core.h>
 #include <nc_server.h>
+#include <stdlib.h>
 
 struct msg *
 rsp_get(struct conn *conn)
@@ -138,6 +139,47 @@ rsp_recv_next(struct context *ctx, struct conn *conn, bool alloc)
     return msg;
 }
 
+static void
+rsp_update_for_scan(struct server_pool *pool,struct msg *request) {
+    struct msg *response = request->peer;
+    struct mbuf *mbuf,*nbuf;
+    struct mbuf *first_mbuf;
+    rstatus_t status;
+    char tmp_str[40];
+    int len;
+
+    for(mbuf=STAILQ_FIRST(&response->mhdr);mbuf!=NULL;mbuf=nbuf){
+        nbuf=STAILQ_NEXT(mbuf,next);
+        if(mbuf_empty(mbuf)) continue;
+
+        first_mbuf=mbuf;
+        break;
+    }
+    ASSERT(nc_strncmp(first_mbuf->pos,"*2\r\n$",strlen("*2\r\n$")) ==0);
+    uint8_t * p=nc_strchr(first_mbuf->pos + sizeof("*2\r\n$"),first_mbuf->last,'\n');
+    unsigned long cursor = strtoul((const char *)(p+1),NULL,10);
+    unsigned long next_cursor;
+    if(cursor == 0 && request->server_index == array_n(&pool->server)-1){
+        // all redis servers have been scanned, and the scan command of the last redis server has returned.
+        return;
+    }else if(cursor ==0 && request->server_index < array_n(&pool->server)-1){
+        // the current redis server have been scanned,now we continue scan next redis server
+        next_cursor=(cursor << NC_MAX_NSERVER_BITS) | (request->server_index+1);
+    }else{
+        // the current redis server scan not finish , go on
+        next_cursor=(cursor << NC_MAX_NSERVER_BITS) | request->server_index;
+    }
+    // discard the old head "*2\r\n$%d\r\n\%dr\n"
+    p=nc_strchr(p+1,first_mbuf->last,'\n');
+    ASSERT(p < first_mbuf->last);
+    first_mbuf->pos = p+1;
+
+    // we get a new head "*2\r\n$%d\r\n\%dr\n", the cursor contain server index
+    len=sprintf(tmp_str,"%ld",next_cursor);
+    status=msg_prepend_format(response,"*2\r\n$%d\r\n%ld\r\n",len,next_cursor);
+    ASSERT(status == NC_OK);
+}
+
 static bool
 rsp_filter(struct context *ctx, struct conn *conn, struct msg *msg)
 {
@@ -261,6 +303,10 @@ rsp_forward(struct context *ctx, struct conn *s_conn, struct msg *msg)
     c_conn = pmsg->owner;
     ASSERT(c_conn->client && !c_conn->proxy);
 
+    if(pmsg->type == MSG_REQ_REDIS_SCAN){
+        rsp_update_for_scan(c_conn->owner,pmsg);
+    }
+
     if (req_done(c_conn, TAILQ_FIRST(&c_conn->omsg_q))) {
         status = event_add_out(ctx->evb, c_conn);
         if (status != NC_OK) {

src/nc_server.c

@@ -700,12 +700,18 @@ server_pool_idx(const struct server_pool *pool, const uint8_t *key, uint32_t key
 }
 
 static struct server *
-server_pool_server(struct server_pool *pool, const uint8_t *key, uint32_t keylen)
+server_pool_server(struct server_pool *pool, struct msg *r, const uint8_t *key, uint32_t keylen)
 {
     struct server *server;
     uint32_t idx;
 
-    idx = server_pool_idx(pool, key, keylen);
+    if (r->type == MSG_REQ_REDIS_SCAN) {
+        idx = r->server_index;
+    }else if (r->type == MSG_REQ_REDIS_SCRIPT) {
+        idx = r->server_index;
+    } else {
+        idx = server_pool_idx(pool, key, keylen);
+    }
     server = array_get(&pool->server, idx);
 
     log_debug(LOG_VERB, "key '%.*s' on dist %d maps to server '%.*s'", keylen,
@@ -715,8 +721,7 @@ server_pool_server(struct server_pool *pool, const uint8_t *key, uint32_t keylen
 }
 
 struct conn *
-server_pool_conn(struct context *ctx, struct server_pool *pool, const uint8_t *key,
-                 uint32_t keylen)
+server_pool_conn(struct context *ctx, struct server_pool *pool, struct msg *msg, const uint8_t *key, uint32_t keylen)
 {
     rstatus_t status;
     struct server *server;
@@ -728,7 +733,7 @@ server_pool_conn(struct context *ctx, struct server_pool *pool, const uint8_t *k
     }
 
     /* from a given {key, keylen} pick a server from pool */
-    server = server_pool_server(pool, key, keylen);
+    server = server_pool_server(pool, msg, key, keylen);
     if (server == NULL) {
         return NULL;
     }

src/nc_server.h

@@ -137,7 +137,7 @@ void server_connected(struct context *ctx, struct conn *conn);
 void server_ok(struct context *ctx, struct conn *conn);
 
 uint32_t server_pool_idx(const struct server_pool *pool, const uint8_t *key, uint32_t keylen);
-struct conn *server_pool_conn(struct context *ctx, struct server_pool *pool, const uint8_t *key, uint32_t keylen);
+struct conn *server_pool_conn(struct context *ctx, struct server_pool *pool, struct msg *msg, const uint8_t *key, uint32_t keylen);
 rstatus_t server_pool_run(struct server_pool *pool);
 rstatus_t server_pool_preconnect(struct context *ctx);
 void server_pool_disconnect(struct context *ctx);