SERVER-49512 Use SBE global environment to store the timezone database and use it with dateFromParts expression

5 files changed, 46 insertions, 30 deletions

diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 63a6c7952aa..5f7991b70b7 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -58,6 +58,13 @@ namespace mongo::stage_builder {
 std::unique_ptr<sbe::RuntimeEnvironment> makeRuntimeEnvironment(
     OperationContext* opCtx, sbe::value::SlotIdGenerator* slotIdGenerator) {
     auto env = std::make_unique<sbe::RuntimeEnvironment>();
+
+    // Register an unowned global timezone database for datetime expression evaluation.
+    env->registerSlot("timeZoneDB"_sd,
+                      sbe::value::TypeTags::timeZoneDB,
+                      sbe::value::bitcastFrom(getTimeZoneDatabase(opCtx)),
+                      false,
+                      slotIdGenerator);
     return env;
 }
 
@@ -323,7 +330,8 @@ std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildProjectionDefault(
                                             std::move(inputStage),
                                             &_slotIdGenerator,
                                             &_frameIdGenerator,
-                                            *_data.resultSlot);
+                                            *_data.resultSlot,
+                                            _data.env);
     _data.resultSlot = slot;
     return std::move(stage);
 }
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 359d06c7f74..46b512aca6a 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -105,13 +105,13 @@ struct ExpressionVisitorContext {
                              sbe::value::FrameIdGenerator* frameIdGenerator,
                              sbe::value::SlotId rootSlot,
                              sbe::value::SlotVector* relevantSlots,
-                             const TimeZoneDatabase* timeZoneDB)
+                             sbe::RuntimeEnvironment* env)
         : traverseStage(std::move(inputStage)),
           slotIdGenerator(slotIdGenerator),
           frameIdGenerator(frameIdGenerator),
           rootSlot(rootSlot),
           relevantSlots(relevantSlots),
-          timeZoneDB(timeZoneDB) {}
+          runtimeEnvironment(env) {}
 
     void ensureArity(size_t arity) {
         invariant(exprs.size() >= arity);
@@ -285,9 +285,7 @@ struct ExpressionVisitorContext {
     // See the comment above the generateExpression() declaration for an explanation of the
     // 'relevantSlots' list.
     sbe::value::SlotVector* relevantSlots;
-
-    // Unowned timezone database needed to evaluate date/time expressions.
-    const TimeZoneDatabase* timeZoneDB;
+    sbe::RuntimeEnvironment* runtimeEnvironment;
 };
 
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateTraverseHelper(
@@ -1255,22 +1253,22 @@ public:
                                     sbe::EPrimBinary::logicOr, std::move(acc), std::move(b));
                             });
 
-        // The builtins need to access the timeZoneDB in order to compute dates from parts. Here
-        // we pass a pointer to the global timeZoneDB. This should not be freed as it's a singleton
-        // and and it's lifetime is tied to the lifetime of the service context.
-        auto unownedTzDB = sbe::value::bitcastFrom(_context->timeZoneDB);
-
-        auto computeDate = sbe::makeE<sbe::EFunction>(
-            isIsoWeekYear ? "datePartsWeekYear" : "dateParts",
-            sbe::makeEs(sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::timeZoneDB, unownedTzDB),
-                        yearRef.clone(),
-                        monthRef.clone(),
-                        dayRef.clone(),
-                        hourRef.clone(),
-                        minRef.clone(),
-                        secRef.clone(),
-                        millisecRef.clone(),
-                        timeZoneRef.clone()));
+        // Invocation of the datePartsWeekYear and dateParts functions depend on a TimeZoneDatabase
+        // for datetime computation. This global object is registered as an unowned value in the
+        // runtime environment so we pass the corresponding slot to the datePartsWeekYear and
+        // dateParts functions as a variable.
+        auto timeZoneDBSlot = _context->runtimeEnvironment->getSlot("timeZoneDB"_sd);
+        auto computeDate =
+            sbe::makeE<sbe::EFunction>(isIsoWeekYear ? "datePartsWeekYear" : "dateParts",
+                                       sbe::makeEs(sbe::makeE<sbe::EVariable>(timeZoneDBSlot),
+                                                   yearRef.clone(),
+                                                   monthRef.clone(),
+                                                   dayRef.clone(),
+                                                   hourRef.clone(),
+                                                   minRef.clone(),
+                                                   secRef.clone(),
+                                                   millisecRef.clone(),
+                                                   timeZoneRef.clone()));
 
         using iterPair_t = std::vector<std::pair<std::unique_ptr<sbe::EExpression>,
                                                  std::unique_ptr<sbe::EExpression>>>::iterator;
@@ -1827,14 +1825,13 @@ generateExpression(OperationContext* opCtx,
                    sbe::value::SlotIdGenerator* slotIdGenerator,
                    sbe::value::FrameIdGenerator* frameIdGenerator,
                    sbe::value::SlotId rootSlot,
+                   sbe::RuntimeEnvironment* env,
                    sbe::value::SlotVector* relevantSlots) {
     auto tempRelevantSlots = sbe::makeSV(rootSlot);
     relevantSlots = relevantSlots ? relevantSlots : &tempRelevantSlots;
 
-    auto timeZoneDB = getTimeZoneDatabase(opCtx);
-
     ExpressionVisitorContext context(
-        std::move(stage), slotIdGenerator, frameIdGenerator, rootSlot, relevantSlots, timeZoneDB);
+        std::move(stage), slotIdGenerator, frameIdGenerator, rootSlot, relevantSlots, env);
 
     ExpressionPreVisitor preVisitor{&context};
     ExpressionInVisitor inVisitor{&context};
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.h b/src/mongo/db/query/sbe_stage_builder_expression.h
index b0a7049271f..4193abfd847 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.h
+++ b/src/mongo/db/query/sbe_stage_builder_expression.h
@@ -47,7 +47,7 @@ namespace mongo::stage_builder {
  * slots to forward.
  *
  * The 'relevantSlots' is an input/output parameter. Execution of this function may add additional
- * relevant slots to thie list.
+ * relevant slots to the list.
  */
 std::tuple<sbe::value::SlotId, std::unique_ptr<sbe::EExpression>, std::unique_ptr<sbe::PlanStage>>
 generateExpression(OperationContext* opCtx,
@@ -56,6 +56,7 @@ generateExpression(OperationContext* opCtx,
                    sbe::value::SlotIdGenerator* slotIdGenerator,
                    sbe::value::FrameIdGenerator* frameIdGenerator,
                    sbe::value::SlotId inputVar,
+                   sbe::RuntimeEnvironment* env,
                    sbe::value::SlotVector* relevantSlots = nullptr);
 
 /**
diff --git a/src/mongo/db/query/sbe_stage_builder_projection.cpp b/src/mongo/db/query/sbe_stage_builder_projection.cpp
index 37092972125..b80673aac63 100644
--- a/src/mongo/db/query/sbe_stage_builder_projection.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_projection.cpp
@@ -132,6 +132,7 @@ struct ProjectionTraversalVisitorContext {
     // The input stage to this projection and the slot to read a root document from.
     PlanStageType inputStage;
     sbe::value::SlotId inputSlot;
+    sbe::RuntimeEnvironment* env;
     std::stack<NestedLevel> levels;
     std::stack<boost::optional<ProjectEval>> evals;
 
@@ -223,6 +224,7 @@ public:
                                _context->slotIdGenerator,
                                _context->frameIdGenerator,
                                _context->inputSlot,
+                               _context->env,
                                &_context->relevantSlots);
         _context->evals.push({{_context->topLevel().inputSlot, outputSlot, std::move(expr)}});
         _context->topLevel().fieldPathExpressionsTraverseStage = std::move(stage);
@@ -373,9 +375,15 @@ std::pair<sbe::value::SlotId, PlanStageType> generateProjection(
     PlanStageType stage,
     sbe::value::SlotIdGenerator* slotIdGenerator,
     sbe::value::FrameIdGenerator* frameIdGenerator,
-    sbe::value::SlotId inputVar) {
-    ProjectionTraversalVisitorContext context{
-        opCtx, projection->type(), slotIdGenerator, frameIdGenerator, std::move(stage), inputVar};
+    sbe::value::SlotId inputVar,
+    sbe::RuntimeEnvironment* env) {
+    ProjectionTraversalVisitorContext context{opCtx,
+                                              projection->type(),
+                                              slotIdGenerator,
+                                              frameIdGenerator,
+                                              std::move(stage),
+                                              inputVar,
+                                              env};
     context.relevantSlots.push_back(inputVar);
     ProjectionTraversalPreVisitor preVisitor{&context};
     ProjectionTraversalPostVisitor postVisitor{&context};
diff --git a/src/mongo/db/query/sbe_stage_builder_projection.h b/src/mongo/db/query/sbe_stage_builder_projection.h
index eaeac9a34f5..00190b4ccd7 100644
--- a/src/mongo/db/query/sbe_stage_builder_projection.h
+++ b/src/mongo/db/query/sbe_stage_builder_projection.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/db/exec/sbe/expressions/expression.h"
 #include "mongo/db/exec/sbe/stages/stages.h"
 #include "mongo/db/exec/sbe/values/id_generators.h"
 #include "mongo/db/operation_context.h"
@@ -46,6 +47,7 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateProjectio
     std::unique_ptr<sbe::PlanStage> stage,
     sbe::value::SlotIdGenerator* slotIdGenerator,
     sbe::value::FrameIdGenerator* frameIdGenerator,
-    sbe::value::SlotId inputVar);
+    sbe::value::SlotId inputVar,
+    sbe::RuntimeEnvironment* env);
 
 }  // namespace mongo::stage_builder