Enabled -x option (Read last argument from STDIN) on redis-benchmark (#9130)

Add the -x option (Read last argument from STDIN) on redis-benchmark.

Other changes:
To be able to use the code from redis-cli some helper methods were moved to cli_common.(h|c)

Co-authored-by: Oran Agra <oran@redislabs.com>

src/cli_common.c

@@ -28,11 +28,16 @@
  * POSSIBILITY OF SUCH DAMAGE.
  */
 
+#include "fmacros.h"
 #include "cli_common.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <fcntl.h>
 #include <errno.h>
 #include <hiredis.h>
 #include <sdscompat.h> /* Use hiredis' sds compat header that maps sds calls to their hi_ variants */
 #include <sds.h> /* use sds.h from hiredis, so that only one set of sds functions will be present in the binary */
+#include <unistd.h>
 #ifdef USE_OPENSSL
 #include <openssl/ssl.h>
 #include <openssl/err.h>
@@ -192,3 +197,65 @@ int cliSecureInit()
 #endif
     return REDIS_OK;
 }
+
+/* Create an sds from stdin */
+sds readArgFromStdin(void) {
+    char buf[1024];
+    sds arg = sdsempty();
+
+    while(1) {
+        int nread = read(fileno(stdin),buf,1024);
+
+        if (nread == 0) break;
+        else if (nread == -1) {
+            perror("Reading from standard input");
+            exit(1);
+        }
+        arg = sdscatlen(arg,buf,nread);
+    }
+    return arg;
+}
+
+/* Create an sds array from argv, either as-is or by dequoting every
+ * element. When quoted is non-zero, may return a NULL to indicate an
+ * invalid quoted string.
+ *
+ * The caller should free the resulting array of sds strings with
+ * sdsfreesplitres().
+ */
+sds *getSdsArrayFromArgv(int argc,char **argv, int quoted) {
+    sds *res = sds_malloc(sizeof(sds) * argc);
+
+    for (int j = 0; j < argc; j++) {
+        if (quoted) {
+            sds unquoted = unquoteCString(argv[j]);
+            if (!unquoted) {
+                while (--j >= 0) sdsfree(res[j]);
+                sds_free(res);
+                return NULL;
+            }
+            res[j] = unquoted;
+        } else {
+            res[j] = sdsnew(argv[j]);
+        }
+    }
+
+    return res;
+}
+
+/* Unquote a null-terminated string and return it as a binary-safe sds. */
+sds unquoteCString(char *str) {
+    int count;
+    sds *unquoted = sdssplitargs(str, &count);
+    sds res = NULL;
+
+    if (unquoted && count == 1) {
+        res = unquoted[0];
+        unquoted[0] = NULL;
+    }
+
+    if (unquoted)
+        sdsfreesplitres(unquoted, count);
+
+    return res;
+}

src/cli_common.h

@@ -2,6 +2,7 @@
 #define __CLICOMMON_H
 
 #include <hiredis.h>
+#include <sdscompat.h> /* Use hiredis' sds compat header that maps sds calls to their hi_ variants */
 
 typedef struct cliSSLconfig {
     /* Requested SNI, or NULL */
@@ -22,29 +23,16 @@ typedef struct cliSSLconfig {
     char* ciphersuites;
 } cliSSLconfig;
 
-/* Wrapper around redisSecureConnection to avoid hiredis_ssl dependencies if
- * not building with TLS support.
- */
 int cliSecureConnection(redisContext *c, cliSSLconfig config, const char **err);
 
-/* Wrapper around hiredis to allow arbitrary reads and writes.
- *
- * We piggybacks on top of hiredis to achieve transparent TLS support,
- * and use its internal buffers so it can co-exist with commands
- * previously/later issued on the connection.
- *
- * Interface is close to enough to read()/write() so things should mostly
- * work transparently.
- */
-
-/* Write a raw buffer through a redisContext. If we already have something
- * in the buffer (leftovers from hiredis operations) it will be written
- * as well.
- */
 ssize_t cliWriteConn(redisContext *c, const char *buf, size_t buf_len);
 
-/* Wrapper around OpenSSL (libssl and libcrypto) initialisation.
- */
 int cliSecureInit();
 
+sds readArgFromStdin(void);
+
+sds *getSdsArrayFromArgv(int argc,char **argv, int quoted);
+
+sds unquoteCString(char *str);
+
 #endif /* __CLICOMMON_H */

src/redis-benchmark.c

@@ -110,6 +110,7 @@ static struct config {
     int dbnum;
     sds dbnumstr;
     char *tests;
+    int stdinarg; /* get last arg from stdin. (-x option) */
     char *auth;
     const char *user;
     int precision;
@@ -1399,7 +1400,7 @@ static void genBenchmarkRandomData(char *data, int count) {
 }
 
 /* Returns number of consumed options. */
-int parseOptions(int argc, const char **argv) {
+int parseOptions(int argc, char **argv) {
     int i;
     int lastarg;
     int exit_status = 1;
@@ -1430,6 +1431,8 @@ int parseOptions(int argc, const char **argv) {
         } else if (!strcmp(argv[i],"-s")) {
             if (lastarg) goto invalid;
             config.hostsocket = strdup(argv[++i]);
+        } else if (!strcmp(argv[i],"-x")) {
+            config.stdinarg = 1;
         } else if (!strcmp(argv[i],"-a") ) {
             if (lastarg) goto invalid;
             config.auth = strdup(argv[++i]);
@@ -1576,6 +1579,7 @@ usage:
 " -t <tests>         Only run the comma separated list of tests. The test\n"
 "                    names are the same as the ones produced as output.\n"
 " -I                 Idle mode. Just open N idle connections and wait.\n"
+" -x                 Read last argument from STDIN.\n"
 #ifdef USE_OPENSSL
 " --tls              Establish a secure TLS connection.\n"
 " --sni <host>       Server name indication for TLS.\n"
@@ -1673,7 +1677,7 @@ int test_is_selected(char *name) {
     return strstr(config.tests,buf) != NULL;
 }
 
-int main(int argc, const char **argv) {
+int main(int argc, char **argv) {
     int i;
     char *data, *cmd, *tag;
     int len;
@@ -1706,6 +1710,7 @@ int main(int argc, const char **argv) {
     config.hostsocket = NULL;
     config.tests = NULL;
     config.dbnum = 0;
+    config.stdinarg = 0;
     config.auth = NULL;
     config.precision = DEFAULT_LATENCY_PRECISION;
     config.num_threads = 0;
@@ -1812,14 +1817,24 @@ int main(int argc, const char **argv) {
             title = sdscatlen(title, " ", 1);
             title = sdscatlen(title, (char*)argv[i], strlen(argv[i]));
         }
-
+        sds *sds_args = getSdsArrayFromArgv(argc, argv, 0);
+        if (!sds_args) {
+            printf("Invalid quoted string\n");
+            return 1;
+        }
+        if (config.stdinarg) {
+            sds_args = sds_realloc(sds_args,(argc + 1) * sizeof(sds));
+            sds_args[argc] = readArgFromStdin();
+            argc++;
+        }
         do {
-            len = redisFormatCommandArgv(&cmd,argc,argv,NULL);
+            len = redisFormatCommandArgv(&cmd,argc,(const char**)sds_args,NULL);
             // adjust the datasize to the parsed command
             config.datasize = len;
             benchmark(title,cmd,len);
             free(cmd);
         } while(config.loop);
+        sdsfreesplitres(sds_args, argc);
 
         if (config.redis_config != NULL) freeRedisConfig(config.redis_config);
         return 0;

src/redis-cli.c

@@ -775,23 +775,6 @@ static void freeHintsCallback(void *ptr) {
  * Networking / parsing
  *--------------------------------------------------------------------------- */
 
-/* Unquote a null-terminated string and return it as a binary-safe sds. */
-static sds unquoteCString(char *str) {
-    int count;
-    sds *unquoted = sdssplitargs(str, &count);
-    sds res = NULL;
-
-    if (unquoted && count == 1) {
-        res = unquoted[0];
-        unquoted[0] = NULL;
-    }
-
-    if (unquoted)
-        sdsfreesplitres(unquoted, count);
-
-    return res;
-}
-
 /* Send AUTH command to the server */
 static int cliAuth(redisContext *ctx, char *user, char *auth) {
     redisReply *reply;
@@ -1879,23 +1862,6 @@ static void parseEnv() {
     }
 }
 
-static sds readArgFromStdin(void) {
-    char buf[1024];
-    sds arg = sdsempty();
-
-    while(1) {
-        int nread = read(fileno(stdin),buf,1024);
-
-        if (nread == 0) break;
-        else if (nread == -1) {
-            perror("Reading from standard input");
-            exit(1);
-        }
-        arg = sdscatlen(arg,buf,nread);
-    }
-    return arg;
-}
-
 static void usage(int err) {
     sds version = cliVersion();
     FILE *target = err ? stderr: stdout;
@@ -2043,30 +2009,6 @@ static int confirmWithYes(char *msg, int ignore_force) {
     return (nread != 0 && !strcmp("yes", buf));
 }
 
-/* Create an sds array from argv, either as-is or by dequoting every
- * element. When quoted is non-zero, may return a NULL to indicate an
- * invalid quoted string.
- */
-static sds *getSdsArrayFromArgv(int argc, char **argv, int quoted) {
-    sds *res = sds_malloc(sizeof(sds) * argc);
-
-    for (int j = 0; j < argc; j++) {
-        if (quoted) {
-            sds unquoted = unquoteCString(argv[j]);
-            if (!unquoted) {
-                while (--j >= 0) sdsfree(res[j]);
-                sds_free(res);
-                return NULL;
-            }
-            res[j] = unquoted;
-        } else {
-            res[j] = sdsnew(argv[j]);
-        }
-    }
-
-    return res;
-}
-
 static int issueCommandRepeat(int argc, char **argv, long repeat) {
     while (1) {
         if (config.cluster_reissue_command || context == NULL ||