Improve zslRandomLevel

test.diff

diff --git a/src/server.c b/src/server.c
index 16cf285e1..9b4fb30b0 100644
--- a/src/server.c
+++ b/src/server.c
@@ -27,6 +27,7 @@
  * POSSIBILITY OF SUCH DAMAGE.
  */
 
+
 #include "server.h"
 #include "monotonic.h"
 #include "cluster.h"
@@ -6604,7 +6605,7 @@ int iAmMaster(void) {
 
 int __failed_tests = 0;
 int __test_num = 0;
-
+int badgerTest(int argc, char *argv[], int flags);
 /* The flags are the following:
 * --accurate:     Runs tests with more iterations.
 * --large-memory: Enables tests that consume more than 100mb. */
@@ -6625,7 +6626,8 @@ struct redisTest {
     {"zmalloc", zmalloc_test},
     {"sds", sdsTest},
     {"dict", dictTest},
-    {"listpack", listpackTest}
+    {"listpack", listpackTest},
+    {"badger", badgerTest},
 };
 redisTestProc *getTestProcByName(const char *name) {
     int numtests = sizeof(redisTests)/sizeof(struct redisTest);
diff --git a/src/server.h b/src/server.h
index 994e98c32..3a192ae50 100644
--- a/src/server.h
+++ b/src/server.h
@@ -30,6 +30,10 @@
 #ifndef __REDIS_H
 #define __REDIS_H
 
+#ifndef REDIS_TEST
+#define REDIS_TEST
+#endif
+
 #include "fmacros.h"
 #include "config.h"
 #include "solarisfixes.h"
diff --git a/src/t_zset.c b/src/t_zset.c
index e09d4528b..bc0d53d43 100644
--- a/src/t_zset.c
+++ b/src/t_zset.c
@@ -121,11 +121,70 @@ void zslFree(zskiplist *zsl) {
  * levels are less likely to be returned. */
 int zslRandomLevel(void) {
     int level = 1;
+
+    while ((random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF))
+        level += 1;
+    return (level<ZSKIPLIST_MAXLEVEL) ? level : ZSKIPLIST_MAXLEVEL;
+}
+
+static int zslRandomLevel1(void) {
+    int level = 1;
+
     while ((random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF))
         level += 1;
     return (level<ZSKIPLIST_MAXLEVEL) ? level : ZSKIPLIST_MAXLEVEL;
 }
 
+static int zslRandomLevel2(void) {
+    int level = 1;
+    static const long int threshold = ZSKIPLIST_P*RAND_MAX;
+    while (random() < threshold)
+        level += 1;
+    return (level<ZSKIPLIST_MAXLEVEL) ? level : ZSKIPLIST_MAXLEVEL;
+}
+#define L (ZSKIPLIST_MAXLEVEL+1)
+#define start_benchmark() start = timeInMilliseconds()
+#define end_benchmark(msg) do { \
+    elapsed = timeInMilliseconds()-start; \
+    printf(msg ":%lld ms\n", elapsed); \
+} while(0)
+#include <sys/time.h>
+static long long timeInMilliseconds(void) {
+    struct timeval tv;
+
+    gettimeofday(&tv,NULL);
+    return (((long long)tv.tv_sec)*1000)+(tv.tv_usec/1000);
+}
+
+int badgerTest(int argc, char *argv[], int flags) {
+    UNUSED(argc);
+    UNUSED(argv);
+    UNUSED(flags);
+    int i, bins[L];
+    long long start, elapsed;
+    for (i = 0; i < L; i++)
+        bins[i] = 0;
+
+    start_benchmark();
+    for (i = 0; i < 200000000; i++)
+        bins[zslRandomLevel2()]++;
+    end_benchmark("new");
+    for (i = 0; i < L; i++)
+        printf("%02d: %d\n", i, bins[i]);
+
+    for (i = 0; i < L; i++)
+        bins[i] = 0;
+
+    start_benchmark();
+    for (i = 0; i < 200000000; i++)
+        bins[zslRandomLevel1()]++;
+    end_benchmark("orig");
+    for (i = 0; i < L; i++)
+        printf("%02d: %d\n", i, bins[i]);
+
+    return 0;
+}
+
 /* Insert a new node in the skiplist. Assumes the element does not already
  * exist (up to the caller to enforce that). The skiplist takes ownership
  * of the passed SDS string 'ele'. */