Correct result count in Mango execution stats (#867)

Mango execution stats previously incremented the result count
at a point where the final result might be discarded. Instead,
increment the count when we know the result is being included
in the response.

2 files changed, 8 insertions, 6 deletions

diff --git a/src/mango/src/mango_cursor_view.erl b/src/mango/src/mango_cursor_view.erl
index 31e198fca..59dd52226 100644
--- a/src/mango/src/mango_cursor_view.erl
+++ b/src/mango/src/mango_cursor_view.erl
@@ -202,10 +202,7 @@ handle_message({row, Props}, Cursor) ->
                 true ->
                     Cursor2 = update_bookmark_keys(Cursor1, Props),
                     FinalDoc = mango_fields:extract(Doc, Cursor2#cursor.fields),
-                    Cursor3 = Cursor2#cursor {
-                        execution_stats = mango_execution_stats:incr_results_returned(Cursor2#cursor.execution_stats)
-                    },
-                    handle_doc(Cursor3, FinalDoc);
+                    handle_doc(Cursor2, FinalDoc);
                 false ->
                     {ok, Cursor1}
             end;
@@ -230,13 +227,14 @@ handle_all_docs_message(Message, Cursor) ->
 
 handle_doc(#cursor{skip = S} = C, _) when S > 0 ->
     {ok, C#cursor{skip = S - 1}};
-handle_doc(#cursor{limit = L} = C, Doc) when L > 0 ->
+handle_doc(#cursor{limit = L, execution_stats = Stats} = C, Doc) when L > 0 ->
     UserFun = C#cursor.user_fun,
     UserAcc = C#cursor.user_acc,
     {Go, NewAcc} = UserFun({row, Doc}, UserAcc),
     {Go, C#cursor{
         user_acc = NewAcc,
-        limit = L - 1
+        limit = L - 1,
+        execution_stats = mango_execution_stats:incr_results_returned(Stats)
     }};
 handle_doc(C, _Doc) ->
     {stop, C}.
diff --git a/src/mango/test/15-execution-stats-test.py b/src/mango/test/15-execution-stats-test.py
index 67c9e64ec..6b7408b8b 100644
--- a/src/mango/test/15-execution-stats-test.py
+++ b/src/mango/test/15-execution-stats-test.py
@@ -38,6 +38,10 @@ class ExecutionStatsTests(mango.UserDocsTests):
         self.assertEqual(resp["execution_stats"]["results_returned"], 3)
         self.assertGreater(resp["execution_stats"]["execution_time_ms"], 0)
 
+    def test_results_returned_limit(self):
+        resp = self.db.find({"age": {"$lt": 35}}, limit=2, return_raw=True, executionStats=True)
+        self.assertEqual(resp["execution_stats"]["results_returned"], len(resp["docs"]))
+
 @unittest.skipUnless(mango.has_text_service(), "requires text service")
 class ExecutionStatsTests_Text(mango.UserDocsTextTests):