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 )                                                                        \
@@ -300,6 +301,9 @@ struct msg {
     unsigned             fdone:1;         /* all fragments are done? */
     unsigned             swallow:1;       /* swallow response? */
     unsigned             redis:1;         /* redis? */
+
+    uint32_t             scan_server_idx; /* used for store scan redisServer index */
+    uint32_t             max_server_idx;
 };
 
 TAILQ_HEAD(msg_tqh, msg);

src/nc_request.c

@@ -573,7 +573,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, key, keylen, msg);
     if (s_conn == NULL) {
         /*
          * Handle a failure to establish a new connection to a server,

src/nc_server.c

@@ -700,12 +700,37 @@ 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);
+    unsigned long long cursor;
+    unsigned long real_cursor;
+    char arr[16];
+    char format[16];
+
+    if (r->type == MSG_REQ_REDIS_SCAN) {
+        if(keylen == 1 && key[0] == '0') {
+            idx=0;
+        }else{
+            /* If the user request is "scan 45066",
+               the cursor 45066 in the request,
+               we get server_idx=45066 & NC_MAX_NSERVER_MASK=10,
+               real_cursor = 45066>>NC_MAX_NSERVER_BITS = 11,
+               and finally the request sent by the proxy to the redis server will be "scan 00011".
+            */
+            cursor=strtoull(key,NULL,10);
+            idx = cursor & NC_MAX_NSERVER_MASK;
+            real_cursor = (cursor >> NC_MAX_NSERVER_BITS);
+            sprintf(format,"%%0%dd",keylen);
+            sprintf(arr,format,real_cursor);
+            nc_memcpy(key,arr,keylen);
+        }
+        r->scan_server_idx=idx;
+        r->max_server_idx=(pool->server).nelem;
+    }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,
@@ -716,7 +741,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)
+                 uint32_t keylen, struct msg *msg)
 {
     rstatus_t status;
     struct server *server;
@@ -728,7 +753,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, const uint8_t *key, uint32_t keylen, struct msg *msg);
 rstatus_t server_pool_run(struct server_pool *pool);
 rstatus_t server_pool_preconnect(struct context *ctx);
 void server_pool_disconnect(struct context *ctx);

src/proto/nc_redis.c

@@ -283,6 +283,7 @@ redis_argn(const struct msg *r)
     case MSG_REQ_REDIS_GEOSEARCHSTORE:
 
     case MSG_REQ_REDIS_RESTORE:
+    case MSG_REQ_REDIS_SCAN:
         return true;
 
     default:
@@ -750,6 +751,11 @@ redis_parse_req(struct msg *r)
                     break;
                 }
 
+                if (str4icmp(m, 's', 'c', 'a', 'n')) {
+                    r->type = MSG_REQ_REDIS_SCAN;
+                    break;
+                }
+
                 break;
 
             case 5:
@@ -2647,7 +2653,7 @@ redis_pre_coalesce(struct msg *r)
     ASSERT(!r->request);
     ASSERT(pr->request);
 
-    if (pr->frag_id == 0) {
+    if (pr->frag_id == 0 || pr->type == MSG_REQ_REDIS_SCAN) {
         /* do nothing, if not a response to a fragmented request */
         return;
     }
@@ -2945,10 +2951,17 @@ redis_fragment_argx(struct msg *r, uint32_t nserver, struct msg_tqh *frag_msgq,
     return NC_OK;
 }
 
+static rstatus_t redis_fragment_scan(struct msg *r, struct msg_tqh *frag_msgq){
+    r->frag_id = msg_gen_frag_id();
+    r->nfrag = 0;
+    r->frag_owner = r;
+    return NC_OK;
+}
+
 rstatus_t
 redis_fragment(struct msg *r, uint32_t nserver, struct msg_tqh *frag_msgq)
 {
-    if (1 == array_n(r->keys)){
+    if (1 == array_n(r->keys) && r->type != MSG_REQ_REDIS_SCAN){
         return NC_OK;
     }
 
@@ -2963,6 +2976,9 @@ redis_fragment(struct msg *r, uint32_t nserver, struct msg_tqh *frag_msgq)
     case MSG_REQ_REDIS_MSET:
         return redis_fragment_argx(r, nserver, frag_msgq, 2);
 
+    case MSG_REQ_REDIS_SCAN:
+        return redis_fragment_scan(r,frag_msgq);
+
     default:
         return NC_OK;
     }
@@ -3053,6 +3069,46 @@ redis_post_coalesce_mget(struct msg *request)
     }
 }
 
+void redis_post_coalesce_scan(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(strncmp(first_mbuf->pos,"*2\r\n$",strlen("*2\r\n$")) ==0);
+    const char* p=strchr(first_mbuf->pos + sizeof("*2\r\n$"),'\n');
+    unsigned long cursor = strtoul(p+1,NULL,10);
+    unsigned long next_cursor;
+    if(cursor == 0 && request->scan_server_idx == request->max_server_idx-1){
+        // all redis servers have been scanned, and the scan command of the last redis server has returned.
+        return;
+    }else if(cursor ==0 && request->scan_server_idx < request->max_server_idx-1){
+        // the current redis server have been scanned,now we continue scan next redis server
+        next_cursor=(cursor << NC_MAX_NSERVER_BITS) | (request->scan_server_idx+1);
+    }else{
+        // the current redis server scan not finish , go on
+        next_cursor=(cursor << NC_MAX_NSERVER_BITS) | request->scan_server_idx;
+    }
+    // discard the old head "*2\r\n$%d\r\n\%dr\n"
+    p=strchr(p+1,'\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);
+}
+
 /*
  * Post-coalesce handler is invoked when the message is a response to
  * the fragmented multi vector request - 'mget' or 'del' and all the
@@ -3083,6 +3139,9 @@ redis_post_coalesce(struct msg *r)
     case MSG_REQ_REDIS_MSET:
         return redis_post_coalesce_mset(r);
 
+    case MSG_REQ_REDIS_SCAN:
+        return redis_post_coalesce_scan(r);
+
     default:
         NOT_REACHED();
     }

tests/test_redis/test_commands.py

@@ -101,3 +101,27 @@ def test_sscan():
     assert_equal('0', str(cursor))
     assert_equal({b'1'}, set(members))
 
+
+def test_scan():
+    r = getconn()
+    r.set('hello_scan_a1',11)
+    r.set('hello_scan_b1',22)
+    r.hmset("hello_scan_h1",{"a":1,"b":2})
+    r.rpush("hello_scan_l1","a","b")
+    r.sadd("hello_scan_s1","a","a","b")
+    r.zadd("hello_scan_z1",{"one": 1, "two": 2, "three": 3})
+
+    zsetval=r.zrange("hello_scan_z1",0,-1,False,True)
+
+    cursor = 0
+    match_str = "hello_scan_*"
+    rets = []
+    subrets = []
+    while True:
+        cursor,subrets =  r.scan(cursor,match_str,100)
+        if len(subrets):
+            rets.extend(subrets)
+        if cursor == 0:
+            break
+    rets.sort()
+    assert_equal(rets,[b'hello_scan_a1', b'hello_scan_b1', b'hello_scan_h1', b'hello_scan_l1', b'hello_scan_s1', b'hello_scan_z1'])