SERVER-41065 Make agg evaluate() thread safe by passing 'Variables' as a parameter

5 files changed, 270 insertions, 99 deletions

diff --git a/jstests/aggregation/bugs/server9840.js b/jstests/aggregation/bugs/server9840.js
deleted file mode 100644
index 4316730370a..00000000000
--- a/jstests/aggregation/bugs/server9840.js
+++ /dev/null
@@ -1,99 +0,0 @@
-// SERVER-9840 variables in expressions and $let
-
-load('jstests/aggregation/extras/utils.js');
-var t = db.server9840;
-t.drop();
-
-function test(expression, expected) {
-    t.drop();
-    t.insert({zero: 0, one: 1, two: 2, three: 3, nested: {four: 4}});
-
-    // Test in projection:
-    var result = t.aggregate({$project: {_id: 0, res: expression}}).toArray();
-    assert.eq(result, [{res: expected}]);
-
-    // Test in group:
-    var result = t.aggregate({$group: {_id: 0, res: {$sum: expression}}}).toArray();
-    assert.eq(result, [{_id: 0, res: expected}]);
-}
-
-// basics
-test('$two', 2);
-test('$$CURRENT.two', 2);
-test('$$ROOT.two', 2);
-
-// using sub expressions
-test({$add: ['$two', '$$CURRENT.three']}, 5);
-test({$add: ['$$CURRENT.two', '$$ROOT.nested.four']}, 6);
-
-// $let simple
-test({$let: {vars: {a: 10}, in : '$$a'}}, 10);
-test({$let: {vars: {a: '$zero'}, in : '$$a'}}, 0);
-test({$let: {vars: {a: {$add: ['$one', '$two']}, b: 10}, in : {$multiply: ['$$a', '$$b']}}}, 30);
-
-// $let changing CURRENT
-test({$let: {vars: {CURRENT: '$$ROOT.nested'}, in : {$multiply: ['$four', '$$ROOT.two']}}}, 8);
-test({
-    $let: {
-        vars: {CURRENT: '$$CURRENT.nested'},  // using original value of CURRENT
-        in : {$multiply: ['$four', '$$ROOT.two']}
-    }
-},
-     8);
-test({
-    $let: {
-        vars: {CURRENT: '$nested'},  // same as last
-        in : {$multiply: ['$four', '$$ROOT.two']}
-    }
-},
-     8);
-test({
-    $let: {
-        vars: {CURRENT: {$const: {ten: 10}}},  // "artificial" object
-        in : {$multiply: ['$ten', '$$ROOT.two']}
-    }
-},
-     20);
-test({
-    $let: {
-        vars: {CURRENT: '$three'},  // sets current to the number 3 (not an object)
-        in : {$multiply: ['$$CURRENT', '$$ROOT.two']}
-    }
-},
-     6);
-
-// swapping with $let (ensures there is no ordering dependency in vars)
-test({
-    $let: {
-        vars: {x: 6, y: 10},
-        in : {
-            $let: {
-                vars: {x: '$$y', y: '$$x'},  // now {x:10, y:6}
-                in : {$subtract: ['$$x', '$$y']}
-            }
-        }
-    }
-},        // not commutative!
-     4);  // 10-6 not 6-10 or 6-6
-
-// unicode is allowed
-test({$let: {vars: {'日本語': 10}, in : '$$日本語'}}, 10);  // Japanese for "Japanese language"
-
-// Can use ROOT and CURRENT directly with no subfield (SERVER-5916)
-t.drop();
-t.insert({_id: 'obj'});
-assert.eq(t.aggregate({$project: {_id: 0, obj: '$$ROOT'}}).toArray(), [{obj: {_id: 'obj'}}]);
-assert.eq(t.aggregate({$project: {_id: 0, obj: '$$CURRENT'}}).toArray(), [{obj: {_id: 'obj'}}]);
-assert.eq(t.aggregate({$group: {_id: 0, objs: {$push: '$$ROOT'}}}).toArray(),
-          [{_id: 0, objs: [{_id: 'obj'}]}]);
-assert.eq(t.aggregate({$group: {_id: 0, objs: {$push: '$$CURRENT'}}}).toArray(),
-          [{_id: 0, objs: [{_id: 'obj'}]}]);
-
-// check name validity checks
-assertErrorCode(t, {$project: {a: {$let: {vars: {ROOT: 1}, in : '$$ROOT'}}}}, 16867);
-assertErrorCode(t, {$project: {a: {$let: {vars: {FOO: 1}, in : '$$FOO'}}}}, 16867);
-assertErrorCode(t, {$project: {a: {$let: {vars: {_underbar: 1}, in : '$$FOO'}}}}, 16867);
-assertErrorCode(t, {$project: {a: {$let: {vars: {'a.b': 1}, in : '$$FOO'}}}}, 16868);
-assertErrorCode(t, {$project: {a: {$let: {vars: {'a b': 1}, in : '$$FOO'}}}}, 16868);
-assertErrorCode(t, {$project: {a: '$$_underbar'}}, 16870);
-assertErrorCode(t, {$project: {a: '$$with spaces'}}, 16871);
diff --git a/jstests/aggregation/expressions/let.js b/jstests/aggregation/expressions/let.js
new file mode 100644
index 00000000000..5de6db8eebf
--- /dev/null
+++ b/jstests/aggregation/expressions/let.js
@@ -0,0 +1,132 @@
+/**
+ * Basic integration tests for the $let expression.
+ */
+(function() {
+    "use strict";
+
+    load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+
+    let coll = db.agg_expr_let;
+    coll.drop();
+    assert.commandWorked(coll.insert({zero: 0, one: 1, two: 2, three: 3, nested: {four: 4}}));
+
+    function testExpr(expression, output) {
+        const res = coll.aggregate([{$project: {output: expression}}]).toArray();
+        assert.eq(res.length, 1, tojson(res));
+        assert.eq(res[0].output, output, tojson(res));
+
+        // Test in group:
+        const result = coll.aggregate({$group: {_id: 0, res: {$sum: expression}}}).toArray();
+        assert.eq(result, [{_id: 0, res: output}]);
+    }
+
+    // Basic tests.
+    testExpr('$two', 2);
+    testExpr('$$CURRENT.two', 2);
+    testExpr('$$ROOT.two', 2);
+
+    // Using sub expressions.
+    testExpr({$add: ['$two', '$$CURRENT.three']}, 5);
+    testExpr({$add: ['$$CURRENT.two', '$$ROOT.nested.four']}, 6);
+
+    // Verify that the variables defined in $let work.
+    testExpr({$let: {vars: {a: 10}, in : '$$a'}}, 10);
+    testExpr({$let: {vars: {a: '$zero'}, in : '$$a'}}, 0);
+    testExpr({$let: {vars: {a: {$add: ['$one', '$two']}, b: 10}, in : {$multiply: ['$$a', '$$b']}}},
+             30);
+
+    // Verify that the outer level variable works in inner level $let.
+    testExpr({
+        $let: {
+            vars: {var1: 1},
+            in : {$let: {vars: {var2: "$$var1"}, in : {$sum: ["$$var1", "$$var2"]}}}
+        }
+    },
+             2);
+
+    // Verify that the outer level variables get overwritten by inner level variables.
+    testExpr({
+        $let: {
+            vars: {var1: "$one"},
+            in : {$let: {vars: {var2: "$$var1", var1: 3}, in : {$sum: ["$$var2", "$$var1"]}}}
+        }
+    },
+             4);
+
+    // $let changing CURRENT
+    testExpr({$let: {vars: {CURRENT: '$$ROOT.nested'}, in : {$multiply: ['$four', '$$ROOT.two']}}},
+             8);
+    testExpr({
+        $let: {
+            vars: {CURRENT: '$$CURRENT.nested'},  // using original value of CURRENT
+            in : {$multiply: ['$four', '$$ROOT.two']}
+        }
+    },
+             8);
+    testExpr({
+        $let: {
+            vars: {CURRENT: '$nested'},  // same as last
+            in : {$multiply: ['$four', '$$ROOT.two']}
+        }
+    },
+             8);
+    testExpr({
+        $let: {
+            vars: {CURRENT: {$const: {ten: 10}}},  // "artificial" object
+            in : {$multiply: ['$ten', '$$ROOT.two']}
+        }
+    },
+             20);
+    testExpr({
+        $let: {
+            vars: {CURRENT: '$three'},  // sets current to the number 3 (not an object)
+            in : {$multiply: ['$$CURRENT', '$$ROOT.two']}
+        }
+    },
+             6);
+
+    // Swapping with $let (ensures there is no ordering dependency in vars).
+    testExpr({
+        $let: {
+            vars: {x: 6, y: 10},
+            in : {
+                $let: {
+                    vars: {x: '$$y', y: '$$x'},  // now {x:10, y:6}
+                    in : {$subtract: ['$$x', '$$y']}
+                }
+            }
+        }
+    },            // Not commutative!
+             4);  // 10-6 not 6-10 or 6-6
+
+    // Unicode is allowed.
+    testExpr({$let: {vars: {'日本語': 10}, in : '$$日本語'}},
+             10);  // Japanese for "Japanese language".
+
+    // Can use ROOT and CURRENT directly with no subfield (SERVER-5916).
+    coll.drop();
+    coll.insert({_id: 'obj'});
+    assert.eq(coll.aggregate({$project: {_id: 0, obj: '$$ROOT'}}).toArray(), [{obj: {_id: 'obj'}}]);
+    assert.eq(coll.aggregate({$project: {_id: 0, obj: '$$CURRENT'}}).toArray(),
+              [{obj: {_id: 'obj'}}]);
+    assert.eq(coll.aggregate({$group: {_id: 0, objs: {$push: '$$ROOT'}}}).toArray(),
+              [{_id: 0, objs: [{_id: 'obj'}]}]);
+    assert.eq(coll.aggregate({$group: {_id: 0, objs: {$push: '$$CURRENT'}}}).toArray(),
+              [{_id: 0, objs: [{_id: 'obj'}]}]);
+
+    // Check name validity checks.
+    assertErrorCode(coll, {$project: {a: {$let: {vars: {ROOT: 1}, in : '$$ROOT'}}}}, 16867);
+    assertErrorCode(coll, {$project: {a: {$let: {vars: {FOO: 1}, in : '$$FOO'}}}}, 16867);
+    assertErrorCode(coll, {$project: {a: {$let: {vars: {_underbar: 1}, in : '$$FOO'}}}}, 16867);
+    assertErrorCode(coll, {$project: {a: {$let: {vars: {'a.b': 1}, in : '$$FOO'}}}}, 16868);
+    assertErrorCode(coll, {$project: {a: {$let: {vars: {'a b': 1}, in : '$$FOO'}}}}, 16868);
+    assertErrorCode(coll, {$project: {a: '$$_underbar'}}, 16870);
+    assertErrorCode(coll, {$project: {a: '$$with spaces'}}, 16871);
+
+    // Verify that variables defined in '$let' cannot be used to initialize other variables.
+    assertErrorCode(
+        coll,
+        [{$project: {output: {$let: {vars: {var1: "$one", var2: "$$var1"}, in : "$$var1"}}}}],
+        17276);
+
+}());
diff --git a/jstests/aggregation/sources/graphLookup/variables.js b/jstests/aggregation/sources/graphLookup/variables.js
new file mode 100644
index 00000000000..87e2c8b3975
--- /dev/null
+++ b/jstests/aggregation/sources/graphLookup/variables.js
@@ -0,0 +1,41 @@
+/**
+ * Tests to verify that $graphLookup can use the variables defined in an outer scope.
+ */
+(function() {
+    "use strict";
+
+    let local = db.graph_lookup_var_local;
+    let foreign = db.graph_lookup_var_foreign;
+    local.drop();
+    foreign.drop();
+
+    foreign.insert({from: "b", to: "a", _id: 0});
+    local.insert({});
+
+    const basicGraphLookup = {
+        $graphLookup: {
+            from: "graph_lookup_var_foreign",
+            startWith: "$$var1",
+            connectFromField: "from",
+            connectToField: "to",
+            as: "resultsFromGraphLookup"
+        }
+    };
+
+    const lookup = {
+        $lookup: {
+            from: "graph_lookup_var_local",
+            let : {var1: "a"},
+            pipeline: [basicGraphLookup],
+            as: "resultsFromLookup"
+        }
+    };
+
+    // Verify that $graphLookup can use the variable 'var1' which is defined in parent $lookup.
+    let res = local.aggregate([lookup]).toArray();
+    assert.eq(res.length, 1);
+    assert.eq(res[0].resultsFromLookup.length, 1);
+    assert.eq(res[0].resultsFromLookup[0].resultsFromGraphLookup.length, 1);
+    assert.eq(res[0].resultsFromLookup[0].resultsFromGraphLookup[0], {_id: 0, from: "b", to: "a"});
+
+})();
diff --git a/jstests/aggregation/sources/lookup/lookup_subpipeline.js b/jstests/aggregation/sources/lookup/lookup_subpipeline.js
index bb7bddc3eca..39d2ff0d850 100644
--- a/jstests/aggregation/sources/lookup/lookup_subpipeline.js
+++ b/jstests/aggregation/sources/lookup/lookup_subpipeline.js
@@ -376,6 +376,23 @@
                        }
                     }],
                     17276);
+    assertErrorCode(
+        coll,
+        [{$lookup: {let : {var1: 1, var2: "$$var1"}, pipeline: [], from: "from", as: "as"}}],
+        17276);
+    assertErrorCode(coll,
+                    [{
+                       $lookup: {
+                           let : {
+                               var1: {$let: {vars: {var1: 2}, in : "$$var1"}},
+                               var2: {$let: {vars: {var1: 4}, in : "$$var2"}},
+                           },
+                           pipeline: [],
+                           from: "from",
+                           as: "as"
+                       }
+                    }],
+                    17276);
 
     // The dotted path offset of a non-object variable is equivalent referencing an undefined
     // field.
diff --git a/jstests/concurrency/fsm_workloads/schema_validator_with_expr_variables.js b/jstests/concurrency/fsm_workloads/schema_validator_with_expr_variables.js
new file mode 100644
index 00000000000..ea847215a29
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/schema_validator_with_expr_variables.js
@@ -0,0 +1,80 @@
+/**
+ * Test to verify that the schema validator works correctly in a multi-threaded environment, when
+ * $expr uses expressions which mutate variable values while executing ($let, $map etc).
+ * @tags: [requires_non_retryable_writes]
+ */
+
+"use strict";
+
+var $config = (function() {
+    function setup(db, collName) {
+        for (let i = 0; i < 200; ++i) {
+            assertAlways.commandWorked(
+                db[collName].insert({_id: i, a: i, one: 1, counter: 0, array: [0, i]}));
+        }
+
+        // Add a validator which checks that field 'a' has value 5 and sum of the elements in field
+        // 'array' is 5. The expression is purposefully complex so that it can create a stress on
+        // expressions with variables.
+        assertAlways.commandWorked(db.runCommand({
+            collMod: collName,
+            validator: {
+                $expr: {
+                    $and: [
+                        {
+                          $eq: [
+                              5,
+                              {
+                                $let: {
+                                    vars: {item: {$multiply: ["$a", "$one"]}},
+                                    in : {$multiply: ["$$item", "$one"]}
+                                }
+                              }
+                          ]
+                        },
+                        {
+                          $eq: [
+                              5,
+                              {
+                                $sum: {
+                                    $map:
+                                        {"input": "$array", "as": "item", "in": "$$item"}
+                                }
+                              }
+                          ]
+                        }
+                    ]
+                }
+            }
+        }));
+    }
+
+    const states = {
+        applyValidator: function(db, collName) {
+            assertAlways.commandWorked(db[collName].update({_id: 5}, {$inc: {counter: 1}}));
+            assertAlways.commandFailedWithCode(
+                db[collName].update({_id: 4}, {$set: {a: 4}, $inc: {counter: 1}}),
+                ErrorCodes.DocumentValidationFailure);
+
+            // Update all the documents in the collection.
+            assertAlways.commandWorked(db[collName].update(
+                {}, {$set: {a: 5, array: [2, 3]}, $inc: {counter: 1}}, {multi: true}));
+
+            // Validation fails when elements of 'array' doesn't add up to 5.
+            assertAlways.commandFailedWithCode(
+                db[collName].update({_id: 4}, {$set: {a: 5, array: [2, 2]}}),
+                ErrorCodes.DocumentValidationFailure);
+        }
+    };
+
+    let transitions = {applyValidator: {applyValidator: 1}};
+
+    return {
+        threadCount: 50,
+        iterations: 100,
+        states: states,
+        startState: "applyValidator",
+        transitions: transitions,
+        setup: setup
+    };
+})();