SERVER-23336 replace ShardResolver with RemoteCommandTargeter::findHost

12 files changed, 141 insertions, 387 deletions

diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index 38cfbdb9cb3..237810f6c31 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -126,7 +126,6 @@ env.Library(
     source=[
         'chunk_manager_targeter.cpp',
         'cluster_write.cpp',
-        'dbclient_shard_resolver.cpp',
     ],
     LIBDEPS=[
         'client/sharding_client',
diff --git a/src/mongo/s/cluster_write.cpp b/src/mongo/s/cluster_write.cpp
index a2400d59250..4299dbe807c 100644
--- a/src/mongo/s/cluster_write.cpp
+++ b/src/mongo/s/cluster_write.cpp
@@ -44,7 +44,6 @@
 #include "mongo/s/chunk_manager_targeter.h"
 #include "mongo/s/client/dbclient_multi_command.h"
 #include "mongo/s/config.h"
-#include "mongo/s/dbclient_shard_resolver.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/mongos_options.h"
 #include "mongo/s/write_ops/batch_write_exec.h"
@@ -273,9 +272,8 @@ void ClusterWriter::write(OperationContext* txn,
                 return;
             }
 
-            DBClientShardResolver resolver;
             DBClientMultiCommand dispatcher;
-            BatchWriteExec exec(&targeter, &resolver, &dispatcher);
+            BatchWriteExec exec(&targeter, &dispatcher);
             exec.executeBatch(txn, *request, response, &_stats);
         }
 
diff --git a/src/mongo/s/commands/cluster_get_last_error_cmd.cpp b/src/mongo/s/commands/cluster_get_last_error_cmd.cpp
index 25f073a831e..49c4599f121 100644
--- a/src/mongo/s/commands/cluster_get_last_error_cmd.cpp
+++ b/src/mongo/s/commands/cluster_get_last_error_cmd.cpp
@@ -34,8 +34,10 @@
 #include "mongo/db/commands.h"
 #include "mongo/db/lasterror.h"
 #include "mongo/s/client/dbclient_multi_command.h"
+#include "mongo/client/remote_command_targeter.h"
+#include "mongo/s/client/shard_registry.h"
 #include "mongo/s/cluster_last_error_info.h"
-#include "mongo/s/dbclient_shard_resolver.h"
+#include "mongo/s/grid.h"
 #include "mongo/s/write_ops/batch_downconvert.h"
 
 namespace mongo {
@@ -108,11 +110,21 @@ public:
             const ConnectionString& shardEndpoint = it->first;
             const HostOpTime& hot = it->second;
 
-            ConnectionString resolvedHost;
-            status = DBClientShardResolver::findMaster(shardEndpoint, &resolvedHost);
-            if (!status.isOK()) {
+            const ReadPreferenceSetting readPref(ReadPreference::PrimaryOnly, TagSet());
+            auto shard = grid.shardRegistry()->getShard(txn, shardEndpoint.toString());
+            if (!shard) {
+                status =
+                    Status(ErrorCodes::ShardNotFound,
+                           str::stream() << "shard " << shardEndpoint.toString() << " not found");
                 break;
             }
+            auto swHostAndPort = shard->getTargeter()->findHost(readPref);
+            if (!swHostAndPort.isOK()) {
+                status = swHostAndPort.getStatus();
+                break;
+            }
+
+            ConnectionString resolvedHost(swHostAndPort.getValue());
 
             resolvedHostOpTimes[resolvedHost] = hot;
         }
diff --git a/src/mongo/s/commands/cluster_write_cmd.cpp b/src/mongo/s/commands/cluster_write_cmd.cpp
index 470c5c2d789..eafdb4940a8 100644
--- a/src/mongo/s/commands/cluster_write_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_cmd.cpp
@@ -30,6 +30,7 @@
 
 #include "mongo/base/error_codes.h"
 #include "mongo/base/owned_pointer_vector.h"
+#include "mongo/client/remote_command_targeter.h"
 #include "mongo/db/client.h"
 #include "mongo/db/client_basic.h"
 #include "mongo/db/commands.h"
@@ -43,7 +44,6 @@
 #include "mongo/s/cluster_last_error_info.h"
 #include "mongo/s/cluster_write.h"
 #include "mongo/s/commands/cluster_explain.h"
-#include "mongo/s/dbclient_shard_resolver.h"
 #include "mongo/s/grid.h"
 #include "mongo/s/write_ops/batch_upconvert.h"
 #include "mongo/s/write_ops/batched_command_request.h"
@@ -259,7 +259,6 @@ private:
                 return status;
         }
 
-        DBClientShardResolver resolver;
         DBClientMultiCommand dispatcher;
 
         // Assemble requests
@@ -267,11 +266,18 @@ private:
              ++it) {
             const ShardEndpoint* endpoint = *it;
 
-            ConnectionString host;
-            Status status = resolver.chooseWriteHost(txn, endpoint->shardName, &host);
-            if (!status.isOK())
-                return status;
+            const ReadPreferenceSetting readPref(ReadPreference::PrimaryOnly, TagSet());
+            auto shard = grid.shardRegistry()->getShard(txn, endpoint->shardName);
+            if (!shard) {
+                return Status(ErrorCodes::ShardNotFound,
+                              "Could not find shard with id " + endpoint->shardName);
+            }
+            auto swHostAndPort = shard->getTargeter()->findHost(readPref);
+            if (!swHostAndPort.isOK()) {
+                return swHostAndPort.getStatus();
+            }
 
+            ConnectionString host(swHostAndPort.getValue());
             dispatcher.addCommand(host, dbName, command);
         }
 
diff --git a/src/mongo/s/dbclient_shard_resolver.cpp b/src/mongo/s/dbclient_shard_resolver.cpp
deleted file mode 100644
index 3d2a32db338..00000000000
--- a/src/mongo/s/dbclient_shard_resolver.cpp
+++ /dev/null
@@ -1,102 +0,0 @@
-/**
- *    Copyright (C) 2013 MongoDB Inc.
- *
- *    This program is free software: you can redistribute it and/or  modify
- *    it under the terms of the GNU Affero General Public License, version 3,
- *    as published by the Free Software Foundation.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    GNU Affero General Public License for more details.
- *
- *    You should have received a copy of the GNU Affero General Public License
- *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the GNU Affero General Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#include "mongo/platform/basic.h"
-
-#include "mongo/s/dbclient_shard_resolver.h"
-
-#include <set>
-
-#include "mongo/client/replica_set_monitor.h"
-#include "mongo/s/client/shard.h"
-#include "mongo/s/client/shard_registry.h"
-#include "mongo/s/grid.h"
-
-namespace mongo {
-
-using std::string;
-
-Status DBClientShardResolver::chooseWriteHost(OperationContext* txn,
-                                              const string& shardName,
-                                              ConnectionString* shardHost) const {
-    // Internally uses our shard cache, does no reload
-    std::shared_ptr<Shard> shard = grid.shardRegistry()->getShard(txn, shardName);
-    if (!shard) {
-        return Status(ErrorCodes::ShardNotFound,
-                      str::stream() << "unknown shard name " << shardName);
-    }
-
-    return findMaster(shard->getConnString(), shardHost);
-}
-
-Status DBClientShardResolver::findMaster(const ConnectionString& connString,
-                                         ConnectionString* resolvedHost) {
-    if (connString.type() == ConnectionString::MASTER) {
-        *resolvedHost = connString;
-        return Status::OK();
-    }
-
-    dassert(connString.type() == ConnectionString::SET);
-
-    //
-    // If we need to, then get the particular node we're targeting in the replica set
-    //
-
-    // Don't create the monitor unless we need to - fast path
-    ReplicaSetMonitorPtr replMonitor = ReplicaSetMonitor::get(connString.getSetName());
-
-    if (!replMonitor) {
-        // Slow path
-        std::set<HostAndPort> seedServers(connString.getServers().begin(),
-                                          connString.getServers().end());
-        ReplicaSetMonitor::createIfNeeded(connString.getSetName(), seedServers);
-
-        replMonitor = ReplicaSetMonitor::get(connString.getSetName());
-    }
-
-    if (!replMonitor) {
-        return Status(ErrorCodes::ReplicaSetNotFound,
-                      str::stream() << "unknown replica set " << connString.getSetName());
-    }
-
-    try {
-        // This can throw when we don't find a master!
-        HostAndPort masterHostAndPort = replMonitor->getMasterOrUassert();
-        *resolvedHost =
-            fassertStatusOK(28687, ConnectionString::parse(masterHostAndPort.toString()));
-        return Status::OK();
-    } catch (const DBException&) {
-        return Status(ErrorCodes::HostNotFound,
-                      string("could not contact primary for replica set ") +
-                          replMonitor->getName());
-    }
-
-    MONGO_UNREACHABLE;
-}
-
-}  // namespace mongo
diff --git a/src/mongo/s/dbclient_shard_resolver.h b/src/mongo/s/dbclient_shard_resolver.h
deleted file mode 100644
index ba70bdf1dc3..00000000000
--- a/src/mongo/s/dbclient_shard_resolver.h
+++ /dev/null
@@ -1,70 +0,0 @@
-/**
- *    Copyright (C) 2013 MongoDB Inc.
- *
- *    This program is free software: you can redistribute it and/or  modify
- *    it under the terms of the GNU Affero General Public License, version 3,
- *    as published by the Free Software Foundation.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    GNU Affero General Public License for more details.
- *
- *    You should have received a copy of the GNU Affero General Public License
- *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the GNU Affero General Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#pragma once
-
-#include "mongo/s/shard_resolver.h"
-
-namespace mongo {
-
-/**
- * ShardResolver based on the Shard and ReplicaSetMonitor caches.
- *
- * TODO: Currently it's possible for the shard resolver to be stale after we target and remove
- * a shard.  We need to figure out how to refresh.
- */
-class DBClientShardResolver : public ShardResolver {
-public:
-    DBClientShardResolver() {}
-
-    virtual ~DBClientShardResolver() {}
-
-    /**
-     * Returns the current host ConnectionString for a write to a shard.
-     *
-     * Note: Does *not* trigger a refresh of either the shard or replica set monitor caches,
-     * though refreshes may happen unexpectedly between calls.
-     *
-     * Returns ShardNotFound if the shard name is unknown
-     * Returns ReplicaSetNotFound if the replica set is not being tracked
-     * Returns !OK with message if the shard host could not be found for other reasons.
-     */
-    Status chooseWriteHost(OperationContext* txn,
-                           const std::string& shardName,
-                           ConnectionString* shardHost) const override;
-
-    /**
-     * Resolves a replica set connection string to a master or returns an error.
-     *
-     * Returns HostNotFound if the master is not reachable
-     * Returns ReplicaSetNotFound if the replica set is not being tracked
-     */
-    static Status findMaster(const ConnectionString& connString, ConnectionString* resolvedHost);
-};
-
-}  // namespace mongo
diff --git a/src/mongo/s/mock_shard_resolver.h b/src/mongo/s/mock_shard_resolver.h
deleted file mode 100644
index d2263274436..00000000000
--- a/src/mongo/s/mock_shard_resolver.h
+++ /dev/null
@@ -1,53 +0,0 @@
-/**
- *    Copyright (C) 2013 MongoDB Inc.
- *
- *    This program is free software: you can redistribute it and/or  modify
- *    it under the terms of the GNU Affero General Public License, version 3,
- *    as published by the Free Software Foundation.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    GNU Affero General Public License for more details.
- *
- *    You should have received a copy of the GNU Affero General Public License
- *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the GNU Affero General Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#pragma once
-
-#include <string>
-
-#include "mongo/base/status.h"
-#include "mongo/client/dbclientinterface.h"
-#include "mongo/s/shard_resolver.h"
-#include "mongo/unittest/unittest.h"
-
-namespace mongo {
-
-class MockShardResolver : public ShardResolver {
-public:
-    virtual ~MockShardResolver() {}
-
-    Status chooseWriteHost(OperationContext* txn,
-                           const std::string& shardName,
-                           ConnectionString* shardHost) const {
-        *shardHost =
-            unittest::assertGet(ConnectionString::parse(std::string("$") + shardName + ":12345"));
-        return Status::OK();
-    }
-};
-
-}  // namespace mongo
diff --git a/src/mongo/s/shard_resolver.h b/src/mongo/s/shard_resolver.h
deleted file mode 100644
index 4a0fa53f48c..00000000000
--- a/src/mongo/s/shard_resolver.h
+++ /dev/null
@@ -1,61 +0,0 @@
-/**
- *    Copyright (C) 2013 MongoDB Inc.
- *
- *    This program is free software: you can redistribute it and/or  modify
- *    it under the terms of the GNU Affero General Public License, version 3,
- *    as published by the Free Software Foundation.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    GNU Affero General Public License for more details.
- *
- *    You should have received a copy of the GNU Affero General Public License
- *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the GNU Affero General Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#pragma once
-
-#include <string>
-
-#include "mongo/base/status.h"
-#include "mongo/client/dbclientinterface.h"
-#include "mongo/s/shard_resolver.h"
-
-namespace mongo {
-
-class OperationContext;
-
-/**
- * Given a shard name, the ShardResolver resolves a particular host on that shard to contact.
- *
- * TODO: Internally, this is backed by a cache - do we need explicit refresh mechanisms built
- * into this interface?
- */
-class ShardResolver {
-public:
-    virtual ~ShardResolver() {}
-
-    /**
-     * Returns a host we can use for write ops to this shard.
-     *
-     * Returns !OK with message if the shard host could not be found for other reasons.
-     */
-    virtual Status chooseWriteHost(OperationContext* txn,
-                                   const std::string& shardName,
-                                   ConnectionString* shardHost) const = 0;
-};
-
-}  // namespace mongo
diff --git a/src/mongo/s/write_ops/SConscript b/src/mongo/s/write_ops/SConscript
index 7efcb12e935..6cf575fed8c 100644
--- a/src/mongo/s/write_ops/SConscript
+++ b/src/mongo/s/write_ops/SConscript
@@ -33,6 +33,8 @@ env.Library(
     LIBDEPS=[

         'batch_write_types',

         '$BUILD_DIR/mongo/client/connection_string',

+        '$BUILD_DIR/mongo/s/client/sharding_client',

+        '$BUILD_DIR/mongo/s/coreshard',

     ],

 )

 

@@ -74,6 +76,7 @@ env.CppUnitTest(
         'cluster_write_op',

         '$BUILD_DIR/mongo/db/range_arithmetic',

         '$BUILD_DIR/mongo/db/service_context',

+        '$BUILD_DIR/mongo/s/sharding_test_fixture',

     ]

 )

 

@@ -84,6 +87,9 @@ env.CppUnitTest(
         'batch_downconvert_test.cpp',

     ],

     LIBDEPS=[

+        '$BUILD_DIR/mongo/db/auth/authorization_manager_mock_init',

+        '$BUILD_DIR/mongo/db/service_context_noop_init',

+        '$BUILD_DIR/mongo/s/mongoscore',

         'cluster_write_op',

         'cluster_write_op_conversion',

     ]

diff --git a/src/mongo/s/write_ops/batch_write_exec.cpp b/src/mongo/s/write_ops/batch_write_exec.cpp
index e056e05609c..d803fbaf32b 100644
--- a/src/mongo/s/write_ops/batch_write_exec.cpp
+++ b/src/mongo/s/write_ops/batch_write_exec.cpp
@@ -37,7 +37,10 @@
 #include "mongo/base/status.h"
 #include "mongo/bson/util/builder.h"
 #include "mongo/client/connection_string.h"
+#include "mongo/client/remote_command_targeter.h"
 #include "mongo/s/client/multi_command_dispatch.h"
+#include "mongo/s/client/shard_registry.h"
+#include "mongo/s/grid.h"
 #include "mongo/s/write_ops/batch_write_op.h"
 #include "mongo/s/write_ops/write_error_detail.h"
 #include "mongo/util/log.h"
@@ -48,10 +51,8 @@ using std::make_pair;
 using std::stringstream;
 using std::vector;
 
-BatchWriteExec::BatchWriteExec(NSTargeter* targeter,
-                               ShardResolver* resolver,
-                               MultiCommandDispatch* dispatcher)
-    : _targeter(targeter), _resolver(resolver), _dispatcher(dispatcher) {}
+BatchWriteExec::BatchWriteExec(NSTargeter* targeter, MultiCommandDispatch* dispatcher)
+    : _targeter(targeter), _dispatcher(dispatcher) {}
 
 namespace {
 
@@ -169,22 +170,43 @@ void BatchWriteExec::executeBatch(OperationContext* txn,
                     continue;
 
                 // Figure out what host we need to dispatch our targeted batch
-                ConnectionString shardHost;
-                Status resolveStatus =
-                    _resolver->chooseWriteHost(txn, nextBatch->getEndpoint().shardName, &shardHost);
-                if (!resolveStatus.isOK()) {
-                    ++stats->numResolveErrors;
+                bool resolvedHost = true;
+                const ReadPreferenceSetting readPref(ReadPreference::PrimaryOnly, TagSet());
+                auto shard =
+                    grid.shardRegistry()->getShard(txn, nextBatch->getEndpoint().shardName);
+
+                if (!shard) {
+                    Status status = Status(ErrorCodes::ShardNotFound,
+                                           str::stream() << "unknown shard name "
+                                                         << nextBatch->getEndpoint().shardName);
+                    resolvedHost = false;
 
                     // Record a resolve failure
                     // TODO: It may be necessary to refresh the cache if stale, or maybe just
                     // cancel and retarget the batch
                     WriteErrorDetail error;
-                    buildErrorFrom(resolveStatus, &error);
+                    buildErrorFrom(status, &error);
+                    LOG(4) << "unable to send write batch to " << nextBatch->getEndpoint().shardName
+                           << causedBy(status);
+                    batchOp.noteBatchError(*nextBatch, error);
+                }
 
-                    LOG(4) << "unable to send write batch to " << shardHost.toString()
-                           << causedBy(resolveStatus.toString());
+                auto swHostAndPort = shard->getTargeter()->findHost(readPref);
+                if (!swHostAndPort.isOK()) {
+                    resolvedHost = false;
 
+                    // Record a resolve failure
+                    // TODO: It may be necessary to refresh the cache if stale, or maybe just
+                    // cancel and retarget the batch
+                    WriteErrorDetail error;
+                    buildErrorFrom(swHostAndPort.getStatus(), &error);
+                    LOG(4) << "unable to send write batch to " << nextBatch->getEndpoint().shardName
+                           << causedBy(swHostAndPort.getStatus());
                     batchOp.noteBatchError(*nextBatch, error);
+                }
+
+                if (!resolvedHost) {
+                    ++stats->numResolveErrors;
 
                     // We're done with this batch
                     // Clean up when we can't resolve a host
@@ -194,6 +216,8 @@ void BatchWriteExec::executeBatch(OperationContext* txn,
                     continue;
                 }
 
+                ConnectionString shardHost(swHostAndPort.getValue());
+
                 // If we already have a batch for this host, wait until the next time
                 OwnedHostBatchMap::MapType::iterator pendingIt = pendingBatches.find(shardHost);
                 if (pendingIt != pendingBatches.end())
diff --git a/src/mongo/s/write_ops/batch_write_exec.h b/src/mongo/s/write_ops/batch_write_exec.h
index e29a80a9a72..739e16a046d 100644
--- a/src/mongo/s/write_ops/batch_write_exec.h
+++ b/src/mongo/s/write_ops/batch_write_exec.h
@@ -36,7 +36,6 @@
 #include "mongo/bson/timestamp.h"
 #include "mongo/db/repl/optime.h"
 #include "mongo/s/ns_targeter.h"
-#include "mongo/s/shard_resolver.h"
 #include "mongo/s/write_ops/batched_command_request.h"
 #include "mongo/s/write_ops/batched_command_response.h"
 
@@ -65,7 +64,7 @@ class BatchWriteExec {
     MONGO_DISALLOW_COPYING(BatchWriteExec);
 
 public:
-    BatchWriteExec(NSTargeter* targeter, ShardResolver* resolver, MultiCommandDispatch* dispatcher);
+    BatchWriteExec(NSTargeter* targeter, MultiCommandDispatch* dispatcher);
 
     /**
      * Executes a client batch write request by sending child batches to several shard
@@ -83,9 +82,6 @@ private:
     NSTargeter* _targeter;
 
     // Not owned here
-    ShardResolver* _resolver;
-
-    // Not owned here
     MultiCommandDispatch* _dispatcher;
 };
 
diff --git a/src/mongo/s/write_ops/batch_write_exec_test.cpp b/src/mongo/s/write_ops/batch_write_exec_test.cpp
index bf0e4c8685a..74b84e4503c 100644
--- a/src/mongo/s/write_ops/batch_write_exec_test.cpp
+++ b/src/mongo/s/write_ops/batch_write_exec_test.cpp
@@ -28,14 +28,19 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/s/write_ops/batch_write_exec.h"
+
 #include "mongo/base/owned_pointer_vector.h"
-#include "mongo/db/operation_context_noop.h"
+#include "mongo/client/remote_command_targeter_mock.h"
+#include "mongo/client/remote_command_targeter_factory_mock.h"
+#include "mongo/s/catalog/type_shard.h"
 #include "mongo/s/client/mock_multi_write_command.h"
+#include "mongo/s/client/shard_registry.h"
 #include "mongo/s/mock_ns_targeter.h"
-#include "mongo/s/mock_shard_resolver.h"
-#include "mongo/s/write_ops/batch_write_exec.h"
+#include "mongo/s/sharding_test_fixture.h"
 #include "mongo/s/write_ops/batched_command_request.h"
 #include "mongo/s/write_ops/batched_command_response.h"
+#include "mongo/stdx/memory.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -46,35 +51,60 @@ using std::vector;
 
 namespace {
 
+const HostAndPort kTestShardHost = HostAndPort("FakeHost", 12345);
+const HostAndPort kTestConfigShardHost = HostAndPort("FakeConfigHost", 12345);
+const string shardName = "FakeShard";
+
 /**
  * Mimics a single shard backend for a particular collection which can be initialized with a
  * set of write command results to return.
  */
-class MockSingleShardBackend {
+class BatchWriteExecTest : public ShardingTestFixture {
 public:
-    MockSingleShardBackend(OperationContext* txn, const NamespaceString& nss) {
-        // Initialize targeting to a mock shard
-        ShardEndpoint endpoint("shard", ChunkVersion::IGNORED());
+    BatchWriteExecTest() = default;
+    ~BatchWriteExecTest() = default;
+
+    void setUp() override {
+        ShardingTestFixture::setUp();
+        getMessagingPort()->setRemote(HostAndPort("ClientHost", 12345));
+
+        // Set up the RemoteCommandTargeter for the config shard.
+        configTargeter()->setFindHostReturnValue(kTestConfigShardHost);
+
+        // Add a RemoteCommandTargeter for the data shard.
+        std::unique_ptr<RemoteCommandTargeterMock> targeter(
+            stdx::make_unique<RemoteCommandTargeterMock>());
+        targeter->setConnectionStringReturnValue(ConnectionString(kTestShardHost));
+        targeter->setFindHostReturnValue(kTestShardHost);
+        targeterFactory()->addTargeterToReturn(ConnectionString(kTestShardHost),
+                                               std::move(targeter));
+
+        // Set up the shard registry to contain the fake shard.
+        ShardType shardType;
+        shardType.setName(shardName);
+        shardType.setHost(kTestShardHost.toString());
+        std::vector<ShardType> shards{shardType};
+        setupShards(shards);
+
+        // Set up the namespace targeter to target the fake shard.
+        ShardEndpoint endpoint(shardName, ChunkVersion::IGNORED());
         vector<MockRange*> mockRanges;
         mockRanges.push_back(
             new MockRange(endpoint, nss, BSON("x" << MINKEY), BSON("x" << MAXKEY)));
-        targeter.init(mockRanges);
-
-        // Get the connection string for the mock shard
-        resolver.chooseWriteHost(txn, mockRanges.front()->endpoint.shardName, &shardHost);
+        nsTargeter.init(mockRanges);
 
-        // Executor using the mock backend
-        exec.reset(new BatchWriteExec(&targeter, &resolver, &dispatcher));
+        // Make the batch write executor use the mock backend.
+        exec.reset(new BatchWriteExec(&nsTargeter, &dispatcher));
     }
 
     void setMockResults(const vector<MockWriteResult*>& results) {
         dispatcher.init(results);
     }
 
-    ConnectionString shardHost;
+    ConnectionString shardHost{kTestShardHost};
+    NamespaceString nss{"foo.bar"};
 
-    MockNSTargeter targeter;
-    MockShardResolver resolver;
+    MockNSTargeter nsTargeter;
     MockMultiWriteCommand dispatcher;
 
     unique_ptr<BatchWriteExec> exec;
@@ -84,16 +114,11 @@ public:
 // Tests for the BatchWriteExec
 //
 
-TEST(BatchWriteExecTests, SingleOp) {
+TEST_F(BatchWriteExecTest, SingleOp) {
     //
     // Basic execution test
     //
 
-    OperationContextNoop txn;
-    NamespaceString nss("foo.bar");
-
-    MockSingleShardBackend backend(&txn, nss);
-
     BatchedCommandRequest request(BatchedCommandRequest::BatchType_Insert);
     request.setNS(nss);
     request.setOrdered(false);
@@ -103,30 +128,25 @@ TEST(BatchWriteExecTests, SingleOp) {
 
     BatchedCommandResponse response;
     BatchWriteExecStats stats;
-    backend.exec->executeBatch(&txn, request, &response, &stats);
+    exec->executeBatch(operationContext(), request, &response, &stats);
     ASSERT(response.getOk());
 
     ASSERT_EQUALS(stats.numRounds, 1);
 }
 
-TEST(BatchWriteExecTests, SingleOpError) {
+TEST_F(BatchWriteExecTest, SingleOpError) {
     //
     // Basic error test
     //
 
-    OperationContextNoop txn;
-    NamespaceString nss("foo.bar");
-
-    MockSingleShardBackend backend(&txn, nss);
-
     vector<MockWriteResult*> mockResults;
     BatchedCommandResponse errResponse;
     errResponse.setOk(false);
     errResponse.setErrCode(ErrorCodes::UnknownError);
     errResponse.setErrMessage("mock error");
-    mockResults.push_back(new MockWriteResult(backend.shardHost, errResponse));
+    mockResults.push_back(new MockWriteResult(shardHost, errResponse));
 
-    backend.setMockResults(mockResults);
+    setMockResults(mockResults);
 
     BatchedCommandRequest request(BatchedCommandRequest::BatchType_Insert);
     request.setNS(nss);
@@ -137,7 +157,7 @@ TEST(BatchWriteExecTests, SingleOpError) {
 
     BatchedCommandResponse response;
     BatchWriteExecStats stats;
-    backend.exec->executeBatch(&txn, request, &response, &stats);
+    exec->executeBatch(operationContext(), request, &response, &stats);
     ASSERT(response.getOk());
     ASSERT_EQUALS(response.getN(), 0);
     ASSERT(response.isErrDetailsSet());
@@ -152,14 +172,11 @@ TEST(BatchWriteExecTests, SingleOpError) {
 // Test retryable errors
 //
 
-TEST(BatchWriteExecTests, StaleOp) {
+TEST_F(BatchWriteExecTest, StaleOp) {
     //
     // Retry op in exec b/c of stale config
     //
 
-    OperationContextNoop txn;
-    NamespaceString nss("foo.bar");
-
     // Insert request
     BatchedCommandRequest request(BatchedCommandRequest::BatchType_Insert);
     request.setNS(nss);
@@ -168,33 +185,28 @@ TEST(BatchWriteExecTests, StaleOp) {
     // Do single-target, single doc batch write op
     request.getInsertRequest()->addToDocuments(BSON("x" << 1));
 
-    MockSingleShardBackend backend(&txn, nss);
-
     vector<MockWriteResult*> mockResults;
     WriteErrorDetail error;
     error.setErrCode(ErrorCodes::StaleShardVersion);
     error.setErrMessage("mock stale error");
-    mockResults.push_back(new MockWriteResult(backend.shardHost, error));
+    mockResults.push_back(new MockWriteResult(shardHost, error));
 
-    backend.setMockResults(mockResults);
+    setMockResults(mockResults);
 
     // Execute request
     BatchedCommandResponse response;
     BatchWriteExecStats stats;
-    backend.exec->executeBatch(&txn, request, &response, &stats);
+    exec->executeBatch(operationContext(), request, &response, &stats);
     ASSERT(response.getOk());
 
     ASSERT_EQUALS(stats.numStaleBatches, 1);
 }
 
-TEST(BatchWriteExecTests, MultiStaleOp) {
+TEST_F(BatchWriteExecTest, MultiStaleOp) {
     //
     // Retry op in exec multiple times b/c of stale config
     //
 
-    OperationContextNoop txn;
-    NamespaceString nss("foo.bar");
-
     // Insert request
     BatchedCommandRequest request(BatchedCommandRequest::BatchType_Insert);
     request.setNS(nss);
@@ -203,37 +215,32 @@ TEST(BatchWriteExecTests, MultiStaleOp) {
     // Do single-target, single doc batch write op
     request.getInsertRequest()->addToDocuments(BSON("x" << 1));
 
-    MockSingleShardBackend backend(&txn, nss);
-
     vector<MockWriteResult*> mockResults;
     WriteErrorDetail error;
     error.setErrCode(ErrorCodes::StaleShardVersion);
     error.setErrMessage("mock stale error");
     for (int i = 0; i < 3; i++) {
-        mockResults.push_back(new MockWriteResult(backend.shardHost, error));
+        mockResults.push_back(new MockWriteResult(shardHost, error));
     }
 
-    backend.setMockResults(mockResults);
+    setMockResults(mockResults);
 
     // Execute request
     BatchedCommandResponse response;
     BatchWriteExecStats stats;
-    backend.exec->executeBatch(&txn, request, &response, &stats);
+    exec->executeBatch(operationContext(), request, &response, &stats);
     ASSERT(response.getOk());
 
     ASSERT_EQUALS(stats.numStaleBatches, 3);
 }
 
-TEST(BatchWriteExecTests, TooManyStaleOp) {
+TEST_F(BatchWriteExecTest, TooManyStaleOp) {
     //
     // Retry op in exec too many times (without refresh) b/c of stale config
-    // (The mock targeter doesn't report progress on refresh)
+    // (The mock nsTargeter doesn't report progress on refresh)
     // We should report a no progress error for everything in the batch
     //
 
-    OperationContextNoop txn;
-    NamespaceString nss("foo.bar");
-
     // Insert request
     BatchedCommandRequest request(BatchedCommandRequest::BatchType_Insert);
     request.setNS(nss);
@@ -243,23 +250,20 @@ TEST(BatchWriteExecTests, TooManyStaleOp) {
     request.getInsertRequest()->addToDocuments(BSON("x" << 1));
     request.getInsertRequest()->addToDocuments(BSON("x" << 2));
 
-    MockSingleShardBackend backend(&txn, nss);
-
     vector<MockWriteResult*> mockResults;
     WriteErrorDetail error;
     error.setErrCode(ErrorCodes::StaleShardVersion);
     error.setErrMessage("mock stale error");
     for (int i = 0; i < 10; i++) {
-        mockResults.push_back(
-            new MockWriteResult(backend.shardHost, error, request.sizeWriteOps()));
+        mockResults.push_back(new MockWriteResult(shardHost, error, request.sizeWriteOps()));
     }
 
-    backend.setMockResults(mockResults);
+    setMockResults(mockResults);
 
     // Execute request
     BatchedCommandResponse response;
     BatchWriteExecStats stats;
-    backend.exec->executeBatch(&txn, request, &response, &stats);
+    exec->executeBatch(operationContext(), request, &response, &stats);
     ASSERT(response.getOk());
     ASSERT_EQUALS(response.getN(), 0);
     ASSERT(response.isErrDetailsSet());
@@ -267,14 +271,11 @@ TEST(BatchWriteExecTests, TooManyStaleOp) {
     ASSERT_EQUALS(response.getErrDetailsAt(1)->getErrCode(), ErrorCodes::NoProgressMade);
 }
 
-TEST(BatchWriteExecTests, ManyStaleOpWithMigration) {
+TEST_F(BatchWriteExecTest, ManyStaleOpWithMigration) {
     //
     // Retry op in exec many times b/c of stale config, but simulate remote migrations occurring
     //
 
-    OperationContextNoop txn;
-    NamespaceString nss("foo.bar");
-
     // Insert request
     BatchedCommandRequest request(BatchedCommandRequest::BatchType_Insert);
     request.setNS(nss);
@@ -283,22 +284,20 @@ TEST(BatchWriteExecTests, ManyStaleOpWithMigration) {
     // Do single-target, single doc batch write op
     request.getInsertRequest()->addToDocuments(BSON("x" << 1));
 
-    MockSingleShardBackend backend(&txn, nss);
-
     vector<MockWriteResult*> mockResults;
     WriteErrorDetail error;
     error.setErrCode(ErrorCodes::StaleShardVersion);
     error.setErrMessage("mock stale error");
     for (int i = 0; i < 10; i++) {
-        mockResults.push_back(new MockWriteResult(backend.shardHost, error));
+        mockResults.push_back(new MockWriteResult(shardHost, error));
     }
 
-    backend.setMockResults(mockResults);
+    setMockResults(mockResults);
 
     // Execute request
     BatchedCommandResponse response;
     BatchWriteExecStats stats;
-    backend.exec->executeBatch(&txn, request, &response, &stats);
+    exec->executeBatch(operationContext(), request, &response, &stats);
     ASSERT(response.getOk());
 
     ASSERT_EQUALS(stats.numStaleBatches, 6);