SERVER-22793 always clear buffered WorkingSetIDs on saveState

This ensures that the set of WorkingSetIDs does not grow without
bound.

5 files changed, 19 insertions, 18 deletions

diff --git a/src/mongo/db/exec/delete.cpp b/src/mongo/db/exec/delete.cpp
index ac6a57e9076..8a29f214e35 100644
--- a/src/mongo/db/exec/delete.cpp
+++ b/src/mongo/db/exec/delete.cpp
@@ -130,11 +130,7 @@ PlanStage::StageState DeleteStage::work(WorkingSetID* out) {
         // TODO: Do we want to buffer docs and delete them in a group rather than
         // saving/restoring state repeatedly?
         _child->saveState();
-        if (supportsDocLocking()) {
-            // Doc-locking engines require this after saveState() since they don't use
-            // invalidations.
-            WorkingSetCommon::forceFetchAllLocs(_txn, _ws, _collection);
-        }
+        WorkingSetCommon::forceFetchAllLocs(_txn, _ws, _collection);
 
         {
             WriteUnitOfWork wunit(_txn);
diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index 9b94b742fc7..a27c0bbdf7c 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -801,11 +801,7 @@ PlanStage::StageState UpdateStage::work(WorkingSetID* out) {
 
         // Save state before making changes
         _child->saveState();
-        if (supportsDocLocking()) {
-            // Doc-locking engines require this after saveState() since they don't use
-            // invalidations.
-            WorkingSetCommon::forceFetchAllLocs(_txn, _ws, _collection);
-        }
+        WorkingSetCommon::forceFetchAllLocs(_txn, _ws, _collection);
 
         // Do the update and return.
         uint64_t attempt = 1;
diff --git a/src/mongo/db/exec/working_set_common.cpp b/src/mongo/db/exec/working_set_common.cpp
index 82893322c74..8c6c412ba4d 100644
--- a/src/mongo/db/exec/working_set_common.cpp
+++ b/src/mongo/db/exec/working_set_common.cpp
@@ -61,7 +61,13 @@ void WorkingSetCommon::forceFetchAllLocs(OperationContext* txn,
                                          WorkingSet* workingSet,
                                          const Collection* collection) {
     invariant(collection);
-    dassert(supportsDocLocking());
+    if (!supportsDocLocking()) {
+        // Non doc-locking storage engines use invalidations, so we don't need to examine the
+        // buffered working set ids. But we do need to clear the set of ids in order to keep our
+        // memory utilization in check.
+        workingSet->getAndClearIdxIds();
+        return;
+    }
 
     for (auto id : workingSet->getAndClearIdxIds()) {
         if (workingSet->isFree(id)) {
diff --git a/src/mongo/db/exec/working_set_common.h b/src/mongo/db/exec/working_set_common.h
index dff7089f1a6..3cf5ba7a449 100644
--- a/src/mongo/db/exec/working_set_common.h
+++ b/src/mongo/db/exec/working_set_common.h
@@ -44,9 +44,10 @@ public:
                                       const Collection* collection);
 
     /**
-     * Iterates over 'workingSet' members that have transitioned to the LOC_AND_IDX state since
-     * the last yield. For all members still in the LOC_AND_IDX state, fetches the associated
-     * document and puts the member in "loc with unowned obj" state.
+     * For storage engines with document-level concurrency, iterates over 'workingSet' members that
+     * have transitioned to the LOC_AND_IDX state since the last yield. For all members still in the
+     * LOC_AND_IDX state, fetches the associated document and puts the member in "loc with unowned
+     * obj" state.
      *
      * This "force-fetching" is called on saveState() for storage-engines that support document-
      * level locking. This ensures that all WS members are still valid, even after the
diff --git a/src/mongo/db/query/plan_executor.cpp b/src/mongo/db/query/plan_executor.cpp
index 4cf50d9b437..039917ed626 100644
--- a/src/mongo/db/query/plan_executor.cpp
+++ b/src/mongo/db/query/plan_executor.cpp
@@ -247,16 +247,18 @@ void PlanExecutor::saveState() {
         _root->saveState();
     }
 
-    // Doc-locking storage engines drop their transactional context after saving state.
     // The query stages inside this stage tree might buffer record ids (e.g. text, geoNear,
     // mergeSort, sort) which are no longer protected by the storage engine's transactional
     // boundaries. Force-fetch the documents for any such record ids so that we have our
     // own copy in the working set.
     //
-    // This is not necessary for covered plans, as such plans never use buffered record ids
-    // for index or collection lookup.
-    if (supportsDocLocking() && _collection && (!_qs.get() || _qs->root->fetched())) {
+    // This is not necessary for covered plans, as such plans never use buffered record ids for
+    // index or collection lookup.
+    if (_collection && (!_qs.get() || _qs->root->fetched())) {
         WorkingSetCommon::forceFetchAllLocs(_opCtx, _workingSet.get(), _collection);
+    } else {
+        // Clear buffered record ids in order to keep our memory footprint in check.
+        _workingSet->getAndClearIdxIds();
     }
 
     _opCtx = NULL;