SERVER-20869 Add support for probabilistically discarding messages.

5 files changed, 88 insertions, 10 deletions

diff --git a/src/mongo/tools/bridge.cpp b/src/mongo/tools/bridge.cpp
index 2ebaa67f47f..53def3153bd 100644
--- a/src/mongo/tools/bridge.cpp
+++ b/src/mongo/tools/bridge.cpp
@@ -31,6 +31,7 @@
 #include "mongo/platform/basic.h"
 
 #include <boost/optional.hpp>
+#include <cstdint>
 
 #include "mongo/base/init.h"
 #include "mongo/base/initializer.h"
@@ -39,6 +40,7 @@
 #include "mongo/db/service_context.h"
 #include "mongo/db/service_context_noop.h"
 #include "mongo/platform/atomic_word.h"
+#include "mongo/platform/random.h"
 #include "mongo/rpc/command_request.h"
 #include "mongo/rpc/factory.h"
 #include "mongo/rpc/reply_builder_interface.h"
@@ -84,8 +86,11 @@ boost::optional<HostAndPort> extractHostInfo(const rpc::RequestInterface& reques
 
 class Forwarder {
 public:
-    Forwarder(MessagingPort* mp, stdx::mutex* settingsMutex, HostSettingsMap* settings)
-        : _mp(mp), _settingsMutex(settingsMutex), _settings(settings) {}
+    Forwarder(MessagingPort* mp,
+              stdx::mutex* settingsMutex,
+              HostSettingsMap* settings,
+              int64_t seed)
+        : _mp(mp), _settingsMutex(settingsMutex), _settings(settings), _prng(seed) {}
 
     void operator()() {
         DBClientConnection dest;
@@ -181,6 +186,20 @@ public:
                               << ", end connection " << _mp->psock->remoteString();
                         _mp->shutdown();
                         return;
+                    // Forward the message to 'dest' with probability '1 - hostSettings.loss'.
+                    case HostSettings::State::kDiscard:
+                        if (_prng.nextCanonicalDouble() < hostSettings.loss) {
+                            std::string hostName = host ? (host->toString()) : "<unknown>";
+                            if (cmdRequest) {
+                                log() << "Discarding \"" << cmdRequest->getCommandName()
+                                      << "\" command with arguments "
+                                      << cmdRequest->getCommandArgs() << " from " << hostName;
+                            } else {
+                                log() << "Discarding message " << request << " from " << hostName;
+                            }
+                            continue;
+                        }
+                        break;
                 }
 
                 // Send the message we received from '_mp' to 'dest'. 'dest' returns a response for
@@ -288,11 +307,17 @@ private:
 
     stdx::mutex* _settingsMutex;
     HostSettingsMap* _settings;
+
+    PseudoRandom _prng;
 };
 
 class BridgeListener final : public Listener {
 public:
-    BridgeListener() : Listener("bridge", "", mongoBridgeGlobalParams.port) {}
+    BridgeListener()
+        : Listener("bridge", "", mongoBridgeGlobalParams.port),
+          _seedSource(mongoBridgeGlobalParams.seed) {
+        log() << "Setting random seed: " << mongoBridgeGlobalParams.seed;
+    }
 
     void acceptedMP(MessagingPort* mp) final {
         {
@@ -304,7 +329,7 @@ public:
             _ports.insert(mp);
         }
 
-        Forwarder f(mp, &_settingsMutex, &_settings);
+        Forwarder f(mp, &_settingsMutex, &_settings, _seedSource.nextInt64());
         stdx::thread t(f);
         t.detach();
     }
@@ -329,6 +354,8 @@ private:
 
     stdx::mutex _settingsMutex;
     HostSettingsMap _settings;
+
+    PseudoRandom _seedSource;
 };
 
 std::unique_ptr<mongo::BridgeListener> listener;
diff --git a/src/mongo/tools/bridge_commands.cpp b/src/mongo/tools/bridge_commands.cpp
index ad57eebe241..2de947b77f0 100644
--- a/src/mongo/tools/bridge_commands.cpp
+++ b/src/mongo/tools/bridge_commands.cpp
@@ -121,12 +121,51 @@ public:
     }
 };
 
+class CmdDiscardMessagesFrom final : public Command {
+public:
+    Status run(const BSONObj& cmdObj, stdx::mutex* settingsMutex, HostSettingsMap* settings) final {
+        invariant(settingsMutex);
+        invariant(settings);
+
+        std::string hostName;
+        auto status = bsonExtractStringField(cmdObj, kHostFieldName, &hostName);
+        if (!status.isOK()) {
+            return status;
+        }
+
+        double newLoss;
+        auto lossElem = cmdObj["loss"];
+        if (lossElem) {
+            if (!lossElem.isNumber()) {
+                return {ErrorCodes::TypeMismatch, "'loss' field must be a number"};
+            }
+
+            newLoss = lossElem.numberDouble();
+            if (newLoss < 0.0 || newLoss > 1.0) {
+                return {ErrorCodes::BadValue, "'loss' field must be a number between 0 and 1"};
+            }
+        } else {
+            return {ErrorCodes::NoSuchKey, "Missing required field 'loss'"};
+        }
+
+        HostAndPort host(hostName);
+        {
+            stdx::lock_guard<stdx::mutex> lk(*settingsMutex);
+            auto& hostSettings = (*settings)[host];
+            hostSettings.state = HostSettings::State::kDiscard;
+            hostSettings.loss = newLoss;
+        }
+        return Status::OK();
+    }
+};
+
 StringMap<Command*> commandMap;
 
 MONGO_INITIALIZER(RegisterBridgeCommands)(InitializerContext* context) {
     commandMap["delayMessagesFrom"] = new CmdDelayMessagesFrom();
     commandMap["acceptConnectionsFrom"] = new CmdAcceptConnectionsFrom();
     commandMap["rejectConnectionsFrom"] = new CmdRejectConnectionsFrom();
+    commandMap["discardMessagesFrom"] = new CmdDiscardMessagesFrom();
     return Status::OK();
 }
 
diff --git a/src/mongo/tools/bridge_commands.h b/src/mongo/tools/bridge_commands.h
index c29ec0c3f80..5002b697a17 100644
--- a/src/mongo/tools/bridge_commands.h
+++ b/src/mongo/tools/bridge_commands.h
@@ -43,10 +43,11 @@ class StatusWith;
 class StringData;
 
 struct HostSettings {
-    enum class State { kForward, kHangUp };
+    enum class State { kForward, kHangUp, kDiscard };
 
     State state = State::kForward;
     Milliseconds delay{0};
+    double loss = 0.0;
 };
 
 using HostSettingsMap = std::unordered_map<HostAndPort, HostSettings>;
diff --git a/src/mongo/tools/mongobridge_options.cpp b/src/mongo/tools/mongobridge_options.cpp
index 1eb1c3b5a12..005247ee44a 100644
--- a/src/mongo/tools/mongobridge_options.cpp
+++ b/src/mongo/tools/mongobridge_options.cpp
@@ -28,9 +28,11 @@
 
 #include "mongo/tools/mongobridge_options.h"
 
+#include <cstdint>
 #include <iostream>
 
 #include "mongo/base/status.h"
+#include "mongo/platform/random.h"
 #include "mongo/util/options_parser/startup_options.h"
 
 namespace mongo {
@@ -38,19 +40,20 @@ namespace mongo {
 MongoBridgeGlobalParams mongoBridgeGlobalParams;
 
 Status addMongoBridgeOptions(moe::OptionSection* options) {
-    options->addOptionChaining("help", "help", moe::Switch, "produce help message");
+    options->addOptionChaining("help", "help", moe::Switch, "show this usage information");
 
+    options->addOptionChaining("port", "port", moe::Int, "port to listen on for MongoDB messages");
 
-    options->addOptionChaining("port", "port", moe::Int, "port to listen for mongo messages");
+    options->addOptionChaining("seed", "seed", moe::Long, "random seed to use");
 
-
-    options->addOptionChaining("dest", "dest", moe::String, "uri of remote mongod instance");
+    options->addOptionChaining("dest", "dest", moe::String, "URI of remote MongoDB process");
 
     return Status::OK();
 }
 
 void printMongoBridgeHelp(std::ostream* out) {
-    *out << "Usage: mongobridge --port <port> --dest <dest> [ --help ]" << std::endl;
+    *out << "Usage: mongobridge --port <port> --dest <dest> [ --seed <seed> ] [ --help ]"
+         << std::endl;
     *out << moe::startupOptions.helpString();
     *out << std::flush;
 }
@@ -76,6 +79,13 @@ Status storeMongoBridgeOptions(const moe::Environment& params,
     mongoBridgeGlobalParams.port = params["port"].as<int>();
     mongoBridgeGlobalParams.destUri = params["dest"].as<std::string>();
 
+    if (!params.count("seed")) {
+        std::unique_ptr<SecureRandom> seedSource{SecureRandom::create()};
+        mongoBridgeGlobalParams.seed = seedSource->nextInt64();
+    } else {
+        mongoBridgeGlobalParams.seed = params["seed"].as<int64_t>();
+    }
+
     return Status::OK();
 }
 
diff --git a/src/mongo/tools/mongobridge_options.h b/src/mongo/tools/mongobridge_options.h
index 74ce0d082f1..b737156ea86 100644
--- a/src/mongo/tools/mongobridge_options.h
+++ b/src/mongo/tools/mongobridge_options.h
@@ -45,6 +45,7 @@ namespace moe = mongo::optionenvironment;
 
 struct MongoBridgeGlobalParams {
     int port = 0;
+    std::int64_t seed = 0;
     std::string destUri;
 
     MongoBridgeGlobalParams() = default;