SERVER-7408 Correctly handle $skip and $limit in sharded agg

This bug only comes up if the first $skip or $limit precedes the first
$sort or $limit. This is very rare, but should still be handled
correctly.

4 files changed, 51 insertions, 8 deletions

diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index b5ae5c80ef5..efe32173cdb 100755
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -965,7 +965,7 @@ namespace mongo {
 
 
     class DocumentSourceLimit :
-        public DocumentSource {
+        public SplittableDocumentSource {
     public:
         // virtuals from DocumentSource
         virtual ~DocumentSourceLimit();
@@ -988,6 +988,14 @@ namespace mongo {
         static intrusive_ptr<DocumentSourceLimit> create(
             const intrusive_ptr<ExpressionContext> &pExpCtx);
 
+        // Virtuals for SplittableDocumentSource
+        // Need to run on rounter. Running on shard as well is an optimization.
+        virtual intrusive_ptr<DocumentSource> getShardSource() { return this; }
+        virtual intrusive_ptr<DocumentSource> getRouterSource() { return this; }
+
+        long long getLimit() const { return limit; }
+        void setLimit(long long newLimit) { limit = newLimit; }
+
         /**
           Create a limiting DocumentSource from BSON.
 
@@ -1019,7 +1027,7 @@ namespace mongo {
     };
 
     class DocumentSourceSkip :
-        public DocumentSource {
+        public SplittableDocumentSource {
     public:
         // virtuals from DocumentSource
         virtual ~DocumentSourceSkip();
@@ -1042,6 +1050,14 @@ namespace mongo {
         static intrusive_ptr<DocumentSourceSkip> create(
             const intrusive_ptr<ExpressionContext> &pExpCtx);
 
+        // Virtuals for SplittableDocumentSource
+        // Need to run on rounter. Can't run on shards.
+        virtual intrusive_ptr<DocumentSource> getShardSource() { return NULL; }
+        virtual intrusive_ptr<DocumentSource> getRouterSource() { return this; }
+
+        long long getSkip() const { return skip; }
+        void setSkip(long long newSkip) { skip = newSkip; }
+
         /**
           Create a skipping DocumentSource from BSON.
 
diff --git a/src/mongo/db/pipeline/document_source_limit.cpp b/src/mongo/db/pipeline/document_source_limit.cpp
index 422fceb5e98..999d80f0368 100644
--- a/src/mongo/db/pipeline/document_source_limit.cpp
+++ b/src/mongo/db/pipeline/document_source_limit.cpp
@@ -27,9 +27,8 @@
 namespace mongo {
     const char DocumentSourceLimit::limitName[] = "$limit";
 
-    DocumentSourceLimit::DocumentSourceLimit(
-        const intrusive_ptr<ExpressionContext> &pExpCtx):
-        DocumentSource(pExpCtx),
+    DocumentSourceLimit::DocumentSourceLimit(const intrusive_ptr<ExpressionContext> &pExpCtx):
+        SplittableDocumentSource(pExpCtx),
         limit(0),
         count(0) {
     }
diff --git a/src/mongo/db/pipeline/document_source_skip.cpp b/src/mongo/db/pipeline/document_source_skip.cpp
index d4c1fc2caa6..5024b817a90 100644
--- a/src/mongo/db/pipeline/document_source_skip.cpp
+++ b/src/mongo/db/pipeline/document_source_skip.cpp
@@ -28,9 +28,8 @@ namespace mongo {
 
     const char DocumentSourceSkip::skipName[] = "$skip";
 
-    DocumentSourceSkip::DocumentSourceSkip(
-        const intrusive_ptr<ExpressionContext> &pExpCtx):
-        DocumentSource(pExpCtx),
+    DocumentSourceSkip::DocumentSourceSkip(const intrusive_ptr<ExpressionContext> &pExpCtx):
+        SplittableDocumentSource(pExpCtx),
         skip(0),
         count(0) {
     }
diff --git a/src/mongo/db/pipeline/pipeline.cpp b/src/mongo/db/pipeline/pipeline.cpp
index b1ab0c0dbe3..b7ebf86272c 100644
--- a/src/mongo/db/pipeline/pipeline.cpp
+++ b/src/mongo/db/pipeline/pipeline.cpp
@@ -218,6 +218,35 @@ namespace mongo {
             }
         }
 
+        /* Move limits in front of skips. This is more optimal for sharding
+         * since currently, we can only split the pipeline at a single source
+         * and it is better to limit the results coming from each shard
+         */
+        for(int i = pSourceVector->size() - 1; i >= 1 /* not looking at 0 */; i--) {
+            DocumentSourceLimit* limit =
+                dynamic_cast<DocumentSourceLimit*>((*pSourceVector)[i].get());
+            DocumentSourceSkip* skip =
+                dynamic_cast<DocumentSourceSkip*>((*pSourceVector)[i-1].get());
+            if (limit && skip) {
+                // Increase limit by skip since the skipped docs now pass through the $limit
+                limit->setLimit(limit->getLimit() + skip->getSkip());
+                swap((*pSourceVector)[i], (*pSourceVector)[i-1]);
+
+                // Start at back again. This is needed to handle cases with more than 1 $limit
+                // (S means skip, L means limit)
+                //
+                // These two would work without second pass (assuming back to front ordering)
+                // SL   -> LS
+                // SSL  -> LSS
+                //
+                // The following cases need a second pass to handle the second limit
+                // SLL  -> LLS
+                // SSLL -> LLSS
+                // SLSL -> LLSS
+                i = pSourceVector->size(); // decremented before next pass
+            }
+        }
+
         /*
           Coalesce adjacent filters where possible.  Two adjacent filters
           are equivalent to one filter whose predicate is the conjunction of