SERVER-18124 Move 'newest chunk' check to the catalog manager

* Adds ability to specify limit to the chunk listing call.
* Gets rid of the getChunksForShard call since it can be built on getChunks

13 files changed, 86 insertions, 108 deletions

diff --git a/src/mongo/s/catalog/catalog_manager.h b/src/mongo/s/catalog/catalog_manager.h
index bc5e1608dc9..72c805f01eb 100644
--- a/src/mongo/s/catalog/catalog_manager.h
+++ b/src/mongo/s/catalog/catalog_manager.h
@@ -220,17 +220,16 @@ namespace mongo {
                                           std::vector<std::string>* dbs) = 0;
 
         /**
-         * Gets all chunks (of type ChunkType) for a shard.
-         * Returns a !OK status if an error occurs.
-         */
-        virtual Status getChunksForShard(const std::string& shardName,
-                                         std::vector<ChunkType>* chunks) = 0;
-
-        /**
-         * Gets all chunks (of type ChunkType) that satisfy a query.
+         * Gets the requested number of chunks (of type ChunkType) that satisfy a query.
+         *
+         * @param query The query to filter out the results.
+         * @param nToReturn The number of chunk entries to return. 0 means all.
+         * @param chunks Vector entry to receive the results
+         *
          * Returns a !OK status if an error occurs.
          */
         virtual Status getChunks(const Query& query,
+                                 int nToReturn,
                                  std::vector<ChunkType>* chunks) = 0;
 
         /**
diff --git a/src/mongo/s/catalog/legacy/catalog_manager_legacy.cpp b/src/mongo/s/catalog/legacy/catalog_manager_legacy.cpp
index 4037f70d43b..957840dc2a8 100644
--- a/src/mongo/s/catalog/legacy/catalog_manager_legacy.cpp
+++ b/src/mongo/s/catalog/legacy/catalog_manager_legacy.cpp
@@ -1151,17 +1151,12 @@ namespace {
         conn.done();
     }
 
-    Status CatalogManagerLegacy::getChunksForShard(const string& shardName,
-                                                   vector<ChunkType>* chunks) {
-        return getChunks(BSON(ChunkType::shard(shardName)), chunks);
-    }
-
     Status CatalogManagerLegacy::getChunks(const Query& query,
+                                           int nToReturn,
                                            vector<ChunkType>* chunks) {
-        ScopedDbConnection conn(_configServerConnectionString, 30.0);
-        std::unique_ptr<DBClientCursor> cursor(conn->query(ChunkType::ConfigNS,
-                                                           query));
 
+        ScopedDbConnection conn(_configServerConnectionString, 30.0);
+        std::unique_ptr<DBClientCursor> cursor(conn->query(ChunkType::ConfigNS, query, nToReturn));
         if (!cursor.get()) {
             conn.done();
             return Status(ErrorCodes::HostUnreachable, "unable to open chunk cursor");
@@ -1243,10 +1238,10 @@ namespace {
         unsigned failedCount = 0;
         bool sameError = true;
         while (dispatcher.numPending() > 0) {
-            ConnectionString configServer;
+            ConnectionString host;
             RawBSONSerializable responseCmdSerial;
 
-            Status dispatchStatus = dispatcher.recvAny(&configServer,
+            Status dispatchStatus = dispatcher.recvAny(&host,
                                                        &responseCmdSerial);
 
             if (!dispatchStatus.isOK()) {
@@ -1254,7 +1249,7 @@ namespace {
             }
 
             responseObj = responseCmdSerial.toBSON();
-            responses.append(configServer.toString(), responseObj);
+            responses.append(host.toString(), responseObj);
 
             currStatus = Command::getStatusFromCommandResult(responseObj);
             if (!currStatus.isOK()) {
diff --git a/src/mongo/s/catalog/legacy/catalog_manager_legacy.h b/src/mongo/s/catalog/legacy/catalog_manager_legacy.h
index 9be92345517..8d5359a88ce 100644
--- a/src/mongo/s/catalog/legacy/catalog_manager_legacy.h
+++ b/src/mongo/s/catalog/legacy/catalog_manager_legacy.h
@@ -97,10 +97,9 @@ namespace mongo {
         virtual void getDatabasesForShard(const std::string& shardName,
                                           std::vector<std::string>* dbs);
 
-        virtual Status getChunksForShard(const std::string& shardName,
-                                         std::vector<ChunkType>* chunks);
-
-        virtual Status getChunks(const Query& query, std::vector<ChunkType>* chunks);
+        virtual Status getChunks(const Query& query,
+                                 int nToReturn,
+                                 std::vector<ChunkType>* chunks);
 
         virtual Status getAllShards(std::vector<ShardType>* shards);
 
diff --git a/src/mongo/s/chunk_diff-inl.h b/src/mongo/s/chunk_diff-inl.h
index 8ace994416a..995a8c74e2b 100644
--- a/src/mongo/s/chunk_diff-inl.h
+++ b/src/mongo/s/chunk_diff-inl.h
@@ -97,9 +97,8 @@ namespace mongo {
         Query diffQuery = configDiffQuery();
 
         try {
-
             std::vector<ChunkType> chunks;
-            uassertStatusOK(catalogManager->getChunks(diffQuery, &chunks));
+            uassertStatusOK(catalogManager->getChunks(diffQuery, 0, &chunks));
 
             return calculateConfigDiff(chunks);
         }
diff --git a/src/mongo/s/client/shard.cpp b/src/mongo/s/client/shard.cpp
index 1daf5c27a34..5d1a761623a 100644
--- a/src/mongo/s/client/shard.cpp
+++ b/src/mongo/s/client/shard.cpp
@@ -32,13 +32,10 @@
 
 #include "mongo/s/client/shard.h"
 
-#include <boost/make_shared.hpp>
-#include <set>
 #include <string>
 #include <vector>
 
 #include "mongo/client/connpool.h"
-#include "mongo/client/dbclientcursor.h"
 #include "mongo/client/replica_set_monitor.h"
 #include "mongo/db/auth/action_set.h"
 #include "mongo/db/auth/action_type.h"
@@ -51,8 +48,6 @@
 
 namespace mongo {
 
-    using std::list;
-    using std::map;
     using std::string;
     using std::stringstream;
     using std::vector;
@@ -113,7 +108,9 @@ namespace {
           _maxSizeMB(maxSizeMB),
           _isDraining(isDraining) {
 
-        _setAddr(addr);
+        if (!_addr.empty()) {
+            _cs = ConnectionString(addr, ConnectionString::SET);
+        }
     }
 
     Shard::Shard(const std::string& name,
@@ -133,15 +130,11 @@ namespace {
         return shard ? *shard : Shard::EMPTY;
     }
 
-    void Shard::_setAddr( const string& addr ) {
-        _addr = addr;
-        if ( !_addr.empty() ) {
-            _cs = ConnectionString( addr , ConnectionString::SET );
-        }
-    }
+    void Shard::reset(const string& ident) {
+        auto shard = grid.shardRegistry()->findIfExists(ident);
+        invariant(shard);
 
-    void Shard::reset( const string& ident ) {
-        *this = grid.shardRegistry()->findCopy( ident );
+        *this = *shard;
     }
 
     bool Shard::containsNode( const string& node ) const {
diff --git a/src/mongo/s/client/shard.h b/src/mongo/s/client/shard.h
index 72df5387852..3a5652294ff 100644
--- a/src/mongo/s/client/shard.h
+++ b/src/mongo/s/client/shard.h
@@ -174,9 +174,6 @@ namespace mongo {
         static void installShard(const std::string& name, const Shard& shard);
 
     private:
-
-        void _setAddr( const std::string& addr );
-
         std::string    _name;
         std::string    _addr;
         ConnectionString _cs;
diff --git a/src/mongo/s/client/shard_registry.cpp b/src/mongo/s/client/shard_registry.cpp
index d0f734b5697..f811a18cb1e 100644
--- a/src/mongo/s/client/shard_registry.cpp
+++ b/src/mongo/s/client/shard_registry.cpp
@@ -158,15 +158,6 @@ namespace mongo {
         return (i == _rsLookup.end()) ? Shard::EMPTY : *(i->second.get());
     }
 
-    Shard ShardRegistry::findCopy(const string& ident){
-        shared_ptr<Shard> found = _findWithRetry(ident);
-
-        boost::lock_guard<boost::mutex> lk(_mutex);
-        massert(13128, str::stream() << "can't find shard for: " << ident, found.get());
-
-        return *found.get();
-    }
-
     void ShardRegistry::set(const string& name, const Shard& s) {
         shared_ptr<Shard> ss(boost::make_shared<Shard>(s));
 
diff --git a/src/mongo/s/client/shard_registry.h b/src/mongo/s/client/shard_registry.h
index 7925fc2873a..0235322e2c2 100644
--- a/src/mongo/s/client/shard_registry.h
+++ b/src/mongo/s/client/shard_registry.h
@@ -60,11 +60,6 @@ namespace mongo {
          */
         Shard lookupRSName(const std::string& name);
 
-        /**
-         * Useful for ensuring our shard data will not be modified while we use it.
-         */
-        Shard findCopy(const std::string& ident);
-
         void set(const std::string& name, const Shard& s);
 
         void remove(const std::string& name);
diff --git a/src/mongo/s/commands/cluster_kill_op.cpp b/src/mongo/s/commands/cluster_kill_op.cpp
index d007ced2099..cff677983a5 100644
--- a/src/mongo/s/commands/cluster_kill_op.cpp
+++ b/src/mongo/s/commands/cluster_kill_op.cpp
@@ -42,6 +42,8 @@
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/commands.h"
 #include "mongo/s/client/shard.h"
+#include "mongo/s/client/shard_registry.h"
+#include "mongo/s/grid.h"
 #include "mongo/util/log.h"
 #include "mongo/util/mongoutils/str.h"
 
@@ -80,28 +82,35 @@ namespace {
             std::string opToKill;
             uassertStatusOK(bsonExtractStringField(cmdObj, "op", &opToKill));
 
-            auto opSepPos = opToKill.find(':');
+            const auto opSepPos = opToKill.find(':');
 
             uassert(28625,
                     str::stream() << "The op argument to killOp must be of the format shardid:opid"
                                   << " but found \"" << opToKill << '"',
-                    (opToKill.size() >= 3) &&// must have at least N:N
-                    (opSepPos != std::string::npos) &&  // must have ':' as separator
+                    (opToKill.size() >= 3) && // must have at least N:N
+                    (opSepPos != std::string::npos) && // must have ':' as separator
                     (opSepPos != 0) && // can't be :NN
                     (opSepPos != (opToKill.size() - 1))); // can't be NN:
 
             auto shardIdent = opToKill.substr(0, opSepPos);
+            log() << "want to kill op: " << opToKill;
+
+            // Will throw if shard id is not found
+            auto shard = grid.shardRegistry()->findIfExists(shardIdent);
+            if (!shard) {
+                return appendCommandStatus(result,
+                                           Status(ErrorCodes::ShardNotFound,
+                                                  str::stream() << "shard " << shardIdent
+                                                                << " does not exist"));
+            }
+
             auto opId = std::stoi(opToKill.substr(opSepPos + 1));
 
             // shardid is actually the opid - keeping for backwards compatibility.
-            result.append("shard" , shardIdent);
+            result.append("shard", shardIdent);
             result.append("shardid", opId);
 
-            log() << "want to kill op: " << opToKill;
-
-            // Will throw if shard id is not found
-            Shard s(shardIdent);
-            ScopedDbConnection conn(s.getConnString());
+            ScopedDbConnection conn(shard->getConnString());
             BSONObj cmdRes;
             BSONObjBuilder argsBob;
             argsBob.append("op", opId);
@@ -109,6 +118,7 @@ namespace {
             // intentionally ignore return value - that is how legacy killOp worked.
             conn->runPseudoCommand("admin", "killOp", "$cmd.sys.killop", args, cmdRes);
             conn.done();
+
             // The original behavior of killOp on mongos is to always return success, regardless of
             // whether the shard reported success or not.
             return true;
diff --git a/src/mongo/s/commands/cluster_remove_shard_cmd.cpp b/src/mongo/s/commands/cluster_remove_shard_cmd.cpp
index e8e7f217928..7d75cfd70f0 100644
--- a/src/mongo/s/commands/cluster_remove_shard_cmd.cpp
+++ b/src/mongo/s/commands/cluster_remove_shard_cmd.cpp
@@ -136,7 +136,10 @@ namespace {
                 break;
             case ShardDrainingStatus::ONGOING: {
                 vector<ChunkType> chunks;
-                Status status = grid.catalogManager()->getChunksForShard(s.getName(), &chunks);
+                Status status = grid.catalogManager()->getChunks(
+                                            Query(BSON(ChunkType::shard(s.getName()))),
+                                            0,  // return all
+                                            &chunks);
                 if (!status.isOK()) {
                     return appendCommandStatus(result, status);
                 }
diff --git a/src/mongo/s/config.cpp b/src/mongo/s/config.cpp
index a7fde93c451..8b90d530637 100644
--- a/src/mongo/s/config.cpp
+++ b/src/mongo/s/config.cpp
@@ -319,29 +319,31 @@ namespace mongo {
         
         // TODO: We need to keep this first one-chunk check in until we have a more efficient way of
         // creating/reusing a chunk manager, as doing so requires copying the full set of chunks currently
-
-        BSONObj newest;
+        std::vector<ChunkType> newestChunk;
         if ( oldVersion.isSet() && ! forceReload ) {
-            ScopedDbConnection conn(configServer.modelServer(), 30.0);
-            newest = conn->findOne(ChunkType::ConfigNS,
-                                   Query(BSON(ChunkType::ns(ns))).sort(
-                                           ChunkType::DEPRECATED_lastmod(), -1));
-            conn.done();
-            
-            if ( ! newest.isEmpty() ) {
-                ChunkVersion v = ChunkVersion::fromBSON(newest, ChunkType::DEPRECATED_lastmod());
-                if ( v.equals( oldVersion ) ) {
+            uassertStatusOK(grid.catalogManager()->getChunks(
+                                        Query(BSON(ChunkType::ns(ns)))
+                                            .sort(ChunkType::DEPRECATED_lastmod(), -1),
+                                        1,
+                                        &newestChunk));
+
+            if (!newestChunk.empty()) {
+                invariant(newestChunk.size() == 1);
+                ChunkVersion v = newestChunk[0].getVersion();
+                if (v.equals(oldVersion)) {
                     boost::lock_guard<boost::mutex> lk( _lock );
-                    CollectionInfo& ci = _collections[ns];
-                    uassert( 15885 , str::stream() << "not sharded after reloading from chunks : " << ns , ci.isSharded() );
+                    const CollectionInfo& ci = _collections[ns];
+                    uassert(15885,
+                            str::stream() << "not sharded after reloading from chunks : "
+                                          << ns, ci.isSharded());
                     return ci.getCM();
                 }
             }
 
         }
-        else if( ! oldVersion.isSet() ){
-            warning() << "version 0 found when " << ( forceReload ? "reloading" : "checking" ) << " chunk manager"
-                      << ", collection '" << ns << "' initially detected as sharded" << endl;
+        else if (!oldVersion.isSet()) {
+            warning() << "version 0 found when " << (forceReload ? "reloading" : "checking")
+                      << " chunk manager; collection '" << ns << "' initially detected as sharded";
         }
 
         // we are not locked now, and want to load a new ChunkManager
@@ -351,7 +353,7 @@ namespace mongo {
         {
             boost::lock_guard<boost::mutex> lll ( _hitConfigServerLock );
             
-            if ( ! newest.isEmpty() && ! forceReload ) {
+            if (!newestChunk.empty() && !forceReload) {
                 // if we have a target we're going for
                 // see if we've hit already
                 
@@ -359,18 +361,15 @@ namespace mongo {
                 CollectionInfo& ci = _collections[ns];
                 if ( ci.isSharded() && ci.getCM() ) {
 
-                    ChunkVersion currentVersion =
-                        ChunkVersion::fromBSON(newest, ChunkType::DEPRECATED_lastmod());
+                    ChunkVersion currentVersion = newestChunk[0].getVersion();
 
                     // Only reload if the version we found is newer than our own in the same
                     // epoch
-                    if( currentVersion <= ci.getCM()->getVersion() &&
-                        ci.getCM()->getVersion().hasEqualEpoch( currentVersion ) )
-                    {
+                    if (currentVersion <= ci.getCM()->getVersion() &&
+                            ci.getCM()->getVersion().hasEqualEpoch(currentVersion)) {
                         return ci.getCM();
                     }
                 }
-                
             }
             
             temp.reset(new ChunkManager(oldManager->getns(),
diff --git a/src/mongo/s/d_migrate.cpp b/src/mongo/s/d_migrate.cpp
index c443347fe87..b509b5b5784 100644
--- a/src/mongo/s/d_migrate.cpp
+++ b/src/mongo/s/d_migrate.cpp
@@ -1688,31 +1688,34 @@ namespace mongo {
                     warning() << "moveChunk commit outcome ongoing" << migrateLog;
                     sleepsecs( 10 );
 
+                    // look for the chunk in this shard whose version got bumped
+                    // we assume that if that mod made it to the config, the applyOps was successful
                     try {
-                        ScopedDbConnection conn(shardingState.getConfigServer(), 10.0);
-
-                        // look for the chunk in this shard whose version got bumped
-                        // we assume that if that mod made it to the config, the applyOps was successful
-                        BSONObj doc = conn->findOne(ChunkType::ConfigNS,
+                        std::vector<ChunkType> newestChunk;
+                        Status status = grid.catalogManager()->getChunks(
                                                     Query(BSON(ChunkType::ns(ns)))
-                                                        .sort(BSON(ChunkType::DEPRECATED_lastmod() << -1)));
-
+                                                        .sort(ChunkType::DEPRECATED_lastmod(), -1),
+                                                    1,
+                                                    &newestChunk);
+                        uassertStatusOK(status);
+
+                        ChunkVersion checkVersion;
+                        if (!newestChunk.empty()) {
+                            invariant(newestChunk.size() == 1);
+                            checkVersion = newestChunk[0].getVersion();
+                        }
 
-                        ChunkVersion checkVersion(ChunkVersion::fromBSON(doc));
-                        if ( checkVersion.equals( nextVersion ) ) {
+                        if (checkVersion.equals(nextVersion)) {
                             log() << "moveChunk commit confirmed" << migrateLog;
                             errmsg.clear();
 
                         }
                         else {
                             error() << "moveChunk commit failed: version is at "
-                                            << checkVersion << " instead of " << nextVersion << migrateLog;
+                                    << checkVersion << " instead of " << nextVersion << migrateLog;
                             error() << "TERMINATING" << migrateLog;
                             dbexit( EXIT_SHARDING_ERROR );
                         }
-
-                        conn.done();
-
                     }
                     catch ( ... ) {
                         error() << "moveChunk failed to get confirmation of commit" << migrateLog;
diff --git a/src/mongo/s/metadata_loader.cpp b/src/mongo/s/metadata_loader.cpp
index 99529e2eb9e..6b17be2d6dc 100644
--- a/src/mongo/s/metadata_loader.cpp
+++ b/src/mongo/s/metadata_loader.cpp
@@ -34,9 +34,6 @@
 
 #include <vector>
 
-#include "mongo/client/connpool.h"
-#include "mongo/client/dbclientcursor.h"
-#include "mongo/client/dbclientmockcursor.h"
 #include "mongo/s/catalog/catalog_manager.h"
 #include "mongo/s/catalog/type_chunk.h"
 #include "mongo/s/catalog/type_collection.h"
@@ -184,10 +181,8 @@ namespace mongo {
         differ.attach( ns, metadata->_chunksMap, metadata->_collVersion, versionMap );
 
         try {
-
             std::vector<ChunkType> chunks;
-            Status status = catalogManager->getChunks(differ.configDiffQuery(),
-                                                      &chunks);
+            Status status = catalogManager->getChunks(differ.configDiffQuery(), 0, &chunks);
             if (!status.isOK()) {
                 if (status == ErrorCodes::HostUnreachable) {
                     // Make our metadata invalid