SERVER-39183 honor socket disconnect in $where

2 files changed, 6 insertions, 1 deletions

diff --git a/src/mongo/db/commands/mr.cpp b/src/mongo/db/commands/mr.cpp
index 450fe0c53dc..c9867cb05e9 100644
--- a/src/mongo/db/commands/mr.cpp
+++ b/src/mongo/db/commands/mr.cpp
@@ -253,7 +253,6 @@ void JSMapper::map(const BSONObj& o) {
 BSONObj JSFinalizer::finalize(const BSONObj& o) {
     Scope* s = _func.scope();
 
-    Scope::NoDBAccess no = s->disableDBAccess("can't access db inside finalize");
     s->invokeSafe(_func.func(), &o, 0);
 
     // We don't want to use o.objsize() to size b since there are many cases where the point of
diff --git a/src/mongo/db/matcher/expression_where.cpp b/src/mongo/db/matcher/expression_where.cpp
index c6633bc3802..7d1f261370f 100644
--- a/src/mongo/db/matcher/expression_where.cpp
+++ b/src/mongo/db/matcher/expression_where.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/namespace_string.h"
 #include "mongo/scripting/engine.h"
 #include "mongo/stdx/memory.h"
+#include "mongo/util/scopeguard.h"
 
 
 namespace mongo {
@@ -67,6 +68,8 @@ WhereMatchExpression::WhereMatchExpression(OperationContext* opCtx,
         AuthorizationSession::get(Client::getCurrent())->getAuthenticatedUserNamesToken();
 
     _scope = getGlobalScriptEngine()->getPooledScope(_opCtx, _dbName, "where" + userToken);
+    const auto guard = makeGuard([&] { _scope->unregisterOperation(); });
+
     _func = _scope->createFunction(getCode().c_str());
 
     uassert(ErrorCodes::BadValue, "$where compile error", _func);
@@ -76,6 +79,9 @@ bool WhereMatchExpression::matches(const MatchableDocument* doc, MatchDetails* d
     uassert(28692, "$where compile error", _func);
     BSONObj obj = doc->toBSON();
 
+    _scope->registerOperation(Client::getCurrent()->getOperationContext());
+    const auto guard = makeGuard([&] { _scope->unregisterOperation(); });
+
     if (!getScope().isEmpty()) {
         _scope->init(&getScope());
     }