SERVER-20051 check return value of PlanYieldPolicy::yield()

2 files changed, 7 insertions, 4 deletions

diff --git a/src/mongo/db/query/plan_executor.cpp b/src/mongo/db/query/plan_executor.cpp
index ad7e0e76ada..64878802385 100644
--- a/src/mongo/db/query/plan_executor.cpp
+++ b/src/mongo/db/query/plan_executor.cpp
@@ -371,9 +371,11 @@ PlanExecutor::ExecState PlanExecutor::getNextImpl(Snapshotted<BSONObj>* objOut,
         //   3) we need to yield and retry due to a WriteConflictException.
         // In all cases, the actual yielding happens here.
         if (_yieldPolicy->shouldYield()) {
-            _yieldPolicy->yield(fetcher.get());
+            if (!_yieldPolicy->yield(fetcher.get())) {
+                // A return of false from a yield should only happen if we've been killed during the
+                // yield.
+                invariant(killed());
 
-            if (killed()) {
                 if (NULL != objOut) {
                     Status status(ErrorCodes::OperationFailed,
                                   str::stream() << "Operation aborted because: " << *_killReason);
diff --git a/src/mongo/db/query/plan_executor.h b/src/mongo/db/query/plan_executor.h
index ed21f8b047a..4d91e279e6e 100644
--- a/src/mongo/db/query/plan_executor.h
+++ b/src/mongo/db/query/plan_executor.h
@@ -241,9 +241,10 @@ public:
      * Returns true if the state was successfully restored and the execution tree can be
      * work()'d.
      *
-     * If allowed, will yield and retry if a WriteConflictException is encountered.
+     * Returns false if the PlanExecutor was killed while saved. A killed execution tree cannot be
+     * worked and should be deleted.
      *
-     * Returns false otherwise.  The execution tree cannot be worked and should be deleted.
+     * If allowed, will yield and retry if a WriteConflictException is encountered.
      */
     bool restoreState();