limit number of random elements returned based on config, exit loop early if client is being closed

Author: msotheeswaran-sc

src/config.cpp

@@ -2912,6 +2912,7 @@ standardConfig configs[] = {
     createLongLongConfig("repl-backlog-size", NULL, MODIFIABLE_CONFIG, 1, LLONG_MAX, g_pserver->repl_backlog_size, 1024*1024, MEMORY_CONFIG, NULL, updateReplBacklogSize), /* Default: 1mb */
     createLongLongConfig("repl-backlog-disk-reserve", NULL, IMMUTABLE_CONFIG, 0, LLONG_MAX, cserver.repl_backlog_disk_size, 0, MEMORY_CONFIG, NULL, NULL),
     createLongLongConfig("max-snapshot-slip", NULL, MODIFIABLE_CONFIG, 0, 5000, g_pserver->snapshot_slip, 400, 0, NULL, NULL),
+    createLongLongConfig("max-rand-count", NULL, MODIFIABLE_CONFIG, 0, LONG_MAX/2, g_pserver->rand_total_threshold, LONG_MAX/2, 0, NULL, NULL),
 
     /* Unsigned Long Long configs */
     createULongLongConfig("maxmemory", NULL, MODIFIABLE_CONFIG, 0, LLONG_MAX, g_pserver->maxmemory, 0, MEMORY_CONFIG, NULL, updateMaxmemory),

src/server.h

@@ -2729,6 +2729,8 @@ struct redisServer {
     long long repl_batch_offStart = -1;
     long long repl_batch_idxStart = -1;
 
+    long long rand_total_threshold;
+
     int config_soft_shutdown = false;
     bool soft_shutdown = false;
 

src/t_hash.cpp

@@ -1066,6 +1066,8 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
                 addReplyBulkCBuffer(c, key, sdslen(key));
                 if (withvalues)
                     addReplyBulkCBuffer(c, value, sdslen(value));
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
         } else if (hash->encoding == OBJ_ENCODING_ZIPLIST) {
             ziplistEntry *keys, *vals = NULL;
@@ -1079,6 +1081,8 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
                 count -= sample_count;
                 ziplistRandomPairs((unsigned char*)ptrFromObj(hash), sample_count, keys, vals);
                 harndfieldReplyWithZiplist(c, sample_count, keys, vals);
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(keys);
             zfree(vals);
@@ -1231,7 +1235,7 @@ void hrandfieldCommand(client *c) {
             return;
         } else if (c->argc == 4) {
             withvalues = 1;
-            if (l < -LONG_MAX/2 || l > LONG_MAX/2) {
+            if (l < -g_pserver->rand_total_threshold || l > g_pserver->rand_total_threshold) {
                 addReplyError(c,"value is out of range");
                 return;
             }

src/t_set.cpp

@@ -673,6 +673,10 @@ void srandmemberWithCountCommand(client *c) {
     dict *d;
 
     if (getRangeLongFromObjectOrReply(c,c->argv[2],-LONG_MAX,LONG_MAX,&l,NULL) != C_OK) return;
+    if (l < -g_pserver->rand_total_threshold || l > g_pserver->rand_total_threshold) {
+        addReplyError(c,"value is out of range");
+        return;
+    }
     if (l >= 0) {
         count = (unsigned long) l;
     } else {
@@ -706,6 +710,8 @@ void srandmemberWithCountCommand(client *c) {
             } else {
                 addReplyBulkCBuffer(c,ele,sdslen(ele));
             }
+            if (c->flags & CLIENT_CLOSE_ASAP)
+                break;
         }
         return;
     }

src/t_zset.cpp

@@ -4054,6 +4054,8 @@ void zrandmemberWithCountCommand(client *c, long l, int withscores) {
                 addReplyBulkCBuffer(c, key, sdslen(key));
                 if (withscores)
                     addReplyDouble(c, *(double*)dictGetVal(de));
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
         } else if (zsetobj->encoding == OBJ_ENCODING_ZIPLIST) {
             ziplistEntry *keys, *vals = NULL;
@@ -4067,6 +4069,8 @@ void zrandmemberWithCountCommand(client *c, long l, int withscores) {
                 count -= sample_count;
                 ziplistRandomPairs((unsigned char*)ptrFromObj(zsetobj), sample_count, keys, vals);
                 zarndmemberReplyWithZiplist(c, sample_count, keys, vals);
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(keys);
             zfree(vals);
@@ -4216,7 +4220,7 @@ void zrandmemberCommand(client *c) {
             return;
         } else if (c->argc == 4) {
             withscores = 1;
-            if (l < -LONG_MAX/2 || l > LONG_MAX/2) {
+            if (l < -g_pserver->rand_total_threshold || l > g_pserver->rand_total_threshold) {
                 addReplyError(c,"value is out of range");
                 return;
             }

tests/unit/obuf-limits.tcl

@@ -182,4 +182,12 @@ start_server {tags {"obuf-limits"} overrides { server-threads 1 }} {
         assert_equal "v2" [r get k2]
         assert_equal "v3" [r get k3]
     }
+
+    test "Obuf limit, HRANDFIELD with huge count stopped mid-run" {
+        r config set client-output-buffer-limit {normal 1000000 0 0}
+        r hset myhash a b
+        catch {r hrandfield myhash -999999999} e
+        assert_match "*I/O error*" $e
+        reconnect
+    }
 }