SERVER-43641 upgrade random.h

This reverts commit a40b196bd3cecd0b66a6323f57e6f08efe0af392.

3 files changed, 255 insertions, 157 deletions

diff --git a/src/mongo/platform/random.cpp b/src/mongo/platform/random.cpp
index f12ffb57a5d..9ed392adf1c 100644
--- a/src/mongo/platform/random.cpp
+++ b/src/mongo/platform/random.cpp
@@ -39,75 +39,59 @@
 #include <bcrypt.h>
 #else
 #include <errno.h>
+#include <fcntl.h>
 #endif
 
 #define _CRT_RAND_S
+#include <array>
 #include <cstdlib>
 #include <fstream>
 #include <iostream>
 #include <limits>
 #include <memory>
+#include <random>
 
 #include "mongo/util/assert_util.h"
 #include "mongo/util/log.h"
 
-namespace mongo {
-
-// ---- PseudoRandom  -----
+#ifdef _WIN32
+#define SECURE_RANDOM_BCRYPT
+#elif defined(__linux__) || defined(__sun) || defined(__APPLE__) || defined(__FreeBSD__) || \
+    defined(__EMSCRIPTEN__)
+#define SECURE_RANDOM_URANDOM
+#elif defined(__OpenBSD__)
+#define SECURE_RANDOM_ARCFOUR
+#else
+#error "Must implement SecureRandom for platform"
+#endif
 
-uint32_t PseudoRandom::nextUInt32() {
-    uint32_t t = _x ^ (_x << 11);
-    _x = _y;
-    _y = _z;
-    _z = _w;
-    return _w = _w ^ (_w >> 19) ^ (t ^ (t >> 8));
-}
+namespace mongo {
 
 namespace {
-const uint32_t default_y = 362436069;
-const uint32_t default_z = 521288629;
-const uint32_t default_w = 88675123;
-}  // namespace
-
-PseudoRandom::PseudoRandom(uint32_t seed) {
-    _x = seed;
-    _y = default_y;
-    _z = default_z;
-    _w = default_w;
-}
 
-PseudoRandom::PseudoRandom(int32_t seed) : PseudoRandom(static_cast<uint32_t>(seed)) {}
-
-PseudoRandom::PseudoRandom(int64_t seed)
-    : PseudoRandom(static_cast<uint32_t>(seed >> 32) ^ static_cast<uint32_t>(seed)) {}
-
-int32_t PseudoRandom::nextInt32() {
-    return nextUInt32();
-}
-
-int64_t PseudoRandom::nextInt64() {
-    uint64_t a = nextUInt32();
-    uint64_t b = nextUInt32();
-    return (a << 32) | b;
-}
-
-double PseudoRandom::nextCanonicalDouble() {
-    double result;
-    do {
-        auto generated = static_cast<uint64_t>(nextInt64());
-        result = static_cast<double>(generated) / std::numeric_limits<uint64_t>::max();
-    } while (result == 1.0);
-    return result;
-}
-
-// --- SecureRandom ----
+template <size_t N>
+struct Buffer {
+    uint64_t pop() {
+        return arr[--avail];
+    }
+    uint8_t* fillPtr() {
+        return reinterpret_cast<uint8_t*>(arr.data() + avail);
+    }
+    size_t fillSize() {
+        return sizeof(uint64_t) * (arr.size() - avail);
+    }
+    void setFilled() {
+        avail = arr.size();
+    }
 
-SecureRandom::~SecureRandom() {}
+    std::array<uint64_t, N> arr;
+    size_t avail = 0;
+};
 
-#ifdef _WIN32
-class WinSecureRandom : public SecureRandom {
+#if defined(SECURE_RANDOM_BCRYPT)
+class Source {
 public:
-    WinSecureRandom() {
+    Source() {
         auto ntstatus = ::BCryptOpenAlgorithmProvider(
             &_algHandle, BCRYPT_RNG_ALGORITHM, MS_PRIMITIVE_PROVIDER, 0);
         if (ntstatus != STATUS_SUCCESS) {
@@ -118,7 +102,7 @@ public:
         }
     }
 
-    virtual ~WinSecureRandom() {
+    ~Source() {
         auto ntstatus = ::BCryptCloseAlgorithmProvider(_algHandle, 0);
         if (ntstatus != STATUS_SUCCESS) {
             warning() << "Failed to close crypto algorithm provider destroying secure random "
@@ -127,75 +111,99 @@ public:
         }
     }
 
-    int64_t nextInt64() {
-        int64_t value;
-        auto ntstatus =
-            ::BCryptGenRandom(_algHandle, reinterpret_cast<PUCHAR>(&value), sizeof(value), 0);
+    size_t refill(uint8_t* buf, size_t n) {
+        auto ntstatus = ::BCryptGenRandom(_algHandle, reinterpret_cast<PUCHAR>(buf), n, 0);
         if (ntstatus != STATUS_SUCCESS) {
             error() << "Failed to generate random number from secure random object; NTSTATUS: "
                     << ntstatus;
             fassertFailed(28814);
         }
-        return value;
+        return n;
     }
 
 private:
     BCRYPT_ALG_HANDLE _algHandle;
 };
+#endif  // SECURE_RANDOM_BCRYPT
 
-std::unique_ptr<SecureRandom> SecureRandom::create() {
-    return std::make_unique<WinSecureRandom>();
-}
-
-#elif defined(__linux__) || defined(__sun) || defined(__APPLE__) || defined(__FreeBSD__) || \
-    defined(__EMSCRIPTEN__)
-
-class InputStreamSecureRandom : public SecureRandom {
+#if defined(SECURE_RANDOM_URANDOM)
+class Source {
 public:
-    InputStreamSecureRandom(const char* fn) {
-        _in = std::make_unique<std::ifstream>(fn, std::ios::binary | std::ios::in);
-        if (!_in->is_open()) {
-            error() << "cannot open " << fn << " " << strerror(errno);
-            fassertFailed(28839);
-        }
-    }
-
-    int64_t nextInt64() {
-        int64_t r;
-        _in->read(reinterpret_cast<char*>(&r), sizeof(r));
-        if (_in->fail()) {
-            error() << "InputStreamSecureRandom failed to generate random bytes";
-            fassertFailed(28840);
+    size_t refill(uint8_t* buf, size_t n) {
+        size_t i = 0;
+        while (i < n) {
+            ssize_t r;
+            while ((r = read(sharedFd(), buf + i, n - i)) == -1) {
+                if (errno == EINTR) {
+                    continue;
+                } else {
+                    auto errSave = errno;
+                    error() << "SecureRandom: read `" << kFn << "`: " << strerror(errSave);
+                    fassertFailed(28840);
+                }
+            }
+            i += r;
         }
-        return r;
+        return i;
     }
 
 private:
-    std::unique_ptr<std::ifstream> _in;
+    static constexpr const char* kFn = "/dev/urandom";
+    static int sharedFd() {
+        // Retain the urandom fd forever.
+        // Kernel ensures that concurrent `read` calls don't mingle their data.
+        // http://lkml.iu.edu//hypermail/linux/kernel/0412.1/0181.html
+        static const int fd = [] {
+            int f;
+            while ((f = open(kFn, 0)) == -1) {
+                if (errno == EINTR) {
+                    continue;
+                } else {
+                    auto errSave = errno;
+                    error() << "SecureRandom: open `" << kFn << "`: " << strerror(errSave);
+                    fassertFailed(28839);
+                }
+            }
+            return f;
+        }();
+        return fd;
+    }
 };
+#endif  // SECURE_RANDOM_URANDOM
 
-std::unique_ptr<SecureRandom> SecureRandom::create() {
-    return std::make_unique<InputStreamSecureRandom>("/dev/urandom");
-}
+#if defined(SECURE_RANDOM_ARCFOUR)
+class Source {
+public:
+    size_t refill(uint8_t* buf, size_t n) {
+        arc4random_buf(buf, n);
+        return n;
+    }
+};
+#endif  // SECURE_RANDOM_ARCFOUR
 
-#elif defined(__OpenBSD__)
+}  // namespace
 
-class Arc4SecureRandom : public SecureRandom {
+class SecureUrbg::State {
 public:
-    int64_t nextInt64() {
-        int64_t value;
-        arc4random_buf(&value, sizeof(value));
-        return value;
+    uint64_t get() {
+        if (!_buffer.avail) {
+            size_t n = _source.refill(_buffer.fillPtr(), _buffer.fillSize());
+            _buffer.avail += n / sizeof(uint64_t);
+        }
+        return _buffer.pop();
     }
+
+private:
+    Source _source;
+    Buffer<16> _buffer;
 };
 
-std::unique_ptr<SecureRandom> SecureRandom::create() {
-    return std::make_unique<Arc4SecureRandom>();
-}
+SecureUrbg::SecureUrbg() : _state{std::make_unique<State>()} {}
 
-#else
+SecureUrbg::~SecureUrbg() = default;
 
-#error Must implement SecureRandom for platform
+uint64_t SecureUrbg::operator()() {
+    return _state->get();
+}
 
-#endif
 }  // namespace mongo
diff --git a/src/mongo/platform/random.h b/src/mongo/platform/random.h
index 6ef2a9bf114..3521399c896 100644
--- a/src/mongo/platform/random.h
+++ b/src/mongo/platform/random.h
@@ -29,86 +29,142 @@
 
 #pragma once
 
+#include <algorithm>
 #include <cstdint>
+#include <cstring>
 #include <limits>
 #include <memory>
+#include <random>
 
 namespace mongo {
 
 /**
+ * A uniform random bit generator based on XorShift.
  * Uses http://en.wikipedia.org/wiki/Xorshift
  */
-class PseudoRandom {
+class XorShift128 {
 public:
-    PseudoRandom(int32_t seed);
+    using result_type = uint32_t;
 
-    PseudoRandom(uint32_t seed);
+    static constexpr result_type min() {
+        return std::numeric_limits<result_type>::lowest();
+    }
 
-    PseudoRandom(int64_t seed);
+    static constexpr result_type max() {
+        return std::numeric_limits<result_type>::max();
+    }
 
-    int32_t nextInt32();
+    explicit XorShift128(uint32_t seed) : _x{seed} {}
 
-    int64_t nextInt64();
+    result_type operator()() {
+        uint32_t t = _x ^ (_x << 11);
+        _x = _y;
+        _y = _z;
+        _z = _w;
+        return _w = _w ^ (_w >> 19) ^ (t ^ (t >> 8));
+    }
 
-    /**
-     * Returns a random number in the range [0, 1).
-     */
-    double nextCanonicalDouble();
+private:
+    uint32_t _x;  // seed
+    uint32_t _y = 362436069;
+    uint32_t _z = 521288629;
+    uint32_t _w = 88675123;
+};
 
-    /**
-     * @return a number between 0 and max
-     */
-    int32_t nextInt32(int32_t max) {
-        return static_cast<uint32_t>(nextInt32()) % static_cast<uint32_t>(max);
+/** The SecureUrbg impls all produce the full range of uint64_t. */
+class SecureUrbg {
+public:
+    using result_type = uint64_t;
+    static constexpr result_type min() {
+        return std::numeric_limits<result_type>::lowest();
     }
-
-    /**
-     * @return a number between 0 and max
-     */
-    int64_t nextInt64(int64_t max) {
-        return static_cast<uint64_t>(nextInt64()) % static_cast<uint64_t>(max);
+    static constexpr result_type max() {
+        return std::numeric_limits<result_type>::max();
     }
 
-    /**
-     * This returns an object that adapts PseudoRandom such that it
-     * can be used as the third argument to std::shuffle. Note that
-     * the lifetime of the returned object must be a subset of the
-     * lifetime of the PseudoRandom object.
-     */
-    auto urbg() {
+    // Details including State vary by platform and are deferred to the .cpp file.
+    SecureUrbg();
+    ~SecureUrbg();
+    result_type operator()();
 
-        class URBG {
-        public:
-            explicit URBG(PseudoRandom* impl) : _impl(impl) {}
+private:
+    class State;
+    std::unique_ptr<State> _state;
+};
+
+// Provides mongo-traditional functions around a pluggable UniformRandomBitGenerator.
+template <typename Urbg>
+class RandomBase {
+public:
+    using urbg_type = Urbg;
 
-            using result_type = uint64_t;
+    RandomBase() : _urbg{} {}
+    explicit RandomBase(urbg_type u) : _urbg{std::move(u)} {}
+
+    /** The underlying generator */
+    urbg_type& urbg() {
+        return _urbg;
+    }
+
+    /** A random number in the range [0, 1). */
+    double nextCanonicalDouble() {
+        return std::uniform_real_distribution<double>{0, 1}(_urbg);
+    }
 
-            static constexpr result_type min() {
-                return std::numeric_limits<result_type>::min();
-            }
+    /** A number uniformly distributed over all possible values. */
+    int32_t nextInt32() {
+        return _nextAny<int32_t>();
+    }
 
-            static constexpr result_type max() {
-                return std::numeric_limits<result_type>::max();
-            }
+    /** A number uniformly distributed over all possible values. */
+    int64_t nextInt64() {
+        return _nextAny<int64_t>();
+    }
 
-            result_type operator()() {
-                return _impl->nextInt64();
-            }
+    /** A number in the half-open interval [0, max) */
+    int32_t nextInt32(int32_t max) {
+        return std::uniform_int_distribution<int32_t>(0, max - 1)(_urbg);
+    }
 
-        private:
-            PseudoRandom* _impl;
-        };
+    /** A number in the half-open interval [0, max) */
+    int64_t nextInt64(int64_t max) {
+        return std::uniform_int_distribution<int64_t>(0, max - 1)(_urbg);
+    }
 
-        return URBG(this);
+    /** Fill array `buf` with `n` random bytes. */
+    void fill(void* buf, size_t n) {
+        const auto p = static_cast<uint8_t*>(buf);
+        size_t written = 0;
+        while (written < n) {
+            uint64_t t = nextInt64();
+            size_t w = std::min(n - written, sizeof(t));
+            std::memcpy(p + written, &t, w);
+            written += w;
+        }
     }
 
 private:
-    uint32_t nextUInt32();
+    template <typename T>
+    T _nextAny() {
+        using Limits = std::numeric_limits<T>;
+        return std::uniform_int_distribution<T>(Limits::lowest(), Limits::max())(_urbg);
+    }
 
-    uint32_t _x;
-    uint32_t _y;
-    uint32_t _z;
-    uint32_t _w;
+    urbg_type _urbg;
+};
+
+/**
+ * A Pseudorandom generator that's not cryptographically secure, but very fast and small.
+ */
+class PseudoRandom : public RandomBase<XorShift128> {
+    using Base = RandomBase<XorShift128>;
+
+public:
+    explicit PseudoRandom(uint32_t seed) : Base{XorShift128{seed}} {}
+    explicit PseudoRandom(int32_t seed) : PseudoRandom{static_cast<uint32_t>(seed)} {}
+    explicit PseudoRandom(uint64_t seed)
+        : PseudoRandom{static_cast<uint32_t>(seed ^ (seed >> 32))} {}
+    explicit PseudoRandom(int64_t seed) : PseudoRandom{static_cast<uint64_t>(seed)} {}
 };
 
 /**
@@ -116,12 +172,11 @@ private:
  * Suitable for nonce/crypto
  * Slower than PseudoRandom, so only use when really need
  */
-class SecureRandom {
-public:
-    virtual ~SecureRandom();
-
-    virtual int64_t nextInt64() = 0;
+class SecureRandom : public RandomBase<SecureUrbg> {
+    using Base = RandomBase<SecureUrbg>;
 
-    static std::unique_ptr<SecureRandom> create();
+public:
+    using Base::Base;
 };
+
 }  // namespace mongo
diff --git a/src/mongo/platform/random_test.cpp b/src/mongo/platform/random_test.cpp
index ee82a89490f..c3271f8b926 100644
--- a/src/mongo/platform/random_test.cpp
+++ b/src/mongo/platform/random_test.cpp
@@ -134,13 +134,13 @@ TEST(RandomTest, NextCanonicalDistinctValues) {
 }
 
 /**
- * Test that nextCanonicalDouble() always returns values between 0 and 1.
+ * Test that nextCanonicalDouble() is at least very likely to return values in [0,1).
  */
 TEST(RandomTest, NextCanonicalWithinRange) {
     PseudoRandom prng(10);
-    for (int i = 0; i < 100; i++) {
+    for (size_t i = 0; i < 1'000'000; ++i) {
         double next = prng.nextCanonicalDouble();
-        ASSERT_LTE(0.0, next);
+        ASSERT_GTE(next, 0.0);
         ASSERT_LT(next, 1.0);
     }
 }
@@ -211,12 +211,47 @@ TEST(RandomTest, NextInt64InRange) {
     }
 }
 
+/**
+ * Test uniformity of nextInt32(max)
+ */
+TEST(RandomTest, NextInt32Uniformity) {
+    PseudoRandom prng(10);
+    /* Break the range into sections. */
+    /* Check that all sections get roughly equal # of hits */
+    constexpr int32_t kMax = (int32_t{3} << 29) - 1;
+    constexpr size_t kBuckets = 64;
+    constexpr size_t kNIter = 1'000'000;
+    constexpr double mu = kNIter / kBuckets;
+    constexpr double muSqInv = 1. / (mu * mu);
+    std::vector<size_t> hist(kBuckets);
+    for (size_t i = 0; i < kNIter; ++i) {
+        auto next = prng.nextInt32(kMax);
+        ASSERT_GTE(next, 0);
+        ASSERT_LTE(next, kMax);
+        ++hist[double(next) * kBuckets / (kMax + 1)];
+    }
+    if (kDebugBuild) {
+        for (size_t i = 0; i < hist.size(); ++i) {
+            double dev = std::pow(std::pow((hist[i] - mu) / mu, 2), .5);
+            unittest::log() << format(FMT_STRING("  [{:4}] count:{:4}, dev:{:6f}, {}"),
+                                      i,
+                                      hist[i],
+                                      dev,
+                                      std::string(hist[i] / 256, '*'));
+        }
+    }
+    for (size_t i = 0; i < hist.size(); ++i) {
+        double dev = std::pow(std::pow(hist[i] - mu, 2) * muSqInv, .5);
+        ASSERT_LT(dev, 0.1) << format(FMT_STRING("hist[{}]={}, mu={}"), i, hist[i], mu);
+    }
+}
+
 TEST(RandomTest, Secure1) {
-    auto a = SecureRandom::create();
-    auto b = SecureRandom::create();
+    auto a = SecureRandom();
+    auto b = SecureRandom();
 
     for (int i = 0; i < 100; i++) {
-        ASSERT_NOT_EQUALS(a->nextInt64(), b->nextInt64());
+        ASSERT_NOT_EQUALS(a.nextInt64(), b.nextInt64());
     }
 }
 }  // namespace mongo