SERVER-51870 Reduce memory footprint for SBE PlanStages

3 files changed, 79 insertions, 129 deletions

diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index d6305e59c9c..590cd01d804 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -32,7 +32,6 @@ sbeEnv.Library(
         'stages/sort.cpp',
         'stages/sorted_merge.cpp',
         'stages/spool.cpp',
-        'stages/stages.cpp',
         'stages/text_match.cpp',
         'stages/traverse.cpp',
         'stages/union.cpp',
diff --git a/src/mongo/db/exec/sbe/stages/stages.cpp b/src/mongo/db/exec/sbe/stages/stages.cpp
deleted file mode 100644
index fcb436ef7df..00000000000
--- a/src/mongo/db/exec/sbe/stages/stages.cpp
+++ /dev/null
@@ -1,90 +0,0 @@
-/**
- *    Copyright (C) 2019-present MongoDB, Inc.
- *
- *    This program is free software: you can redistribute it and/or modify
- *    it under the terms of the Server Side Public License, version 1,
- *    as published by MongoDB, Inc.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    Server Side Public License for more details.
- *
- *    You should have received a copy of the Server Side Public License
- *    along with this program. If not, see
- *    <http://www.mongodb.com/licensing/server-side-public-license>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the Server Side Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#include "mongo/platform/basic.h"
-
-#include "mongo/db/exec/sbe/stages/stages.h"
-
-#include "mongo/db/operation_context.h"
-
-namespace mongo {
-namespace sbe {
-void CanSwitchOperationContext::detachFromOperationContext() {
-    invariant(_opCtx);
-
-    for (auto&& child : _stage->_children) {
-        child->detachFromOperationContext();
-    }
-
-    doDetachFromOperationContext();
-    _opCtx = nullptr;
-}
-
-void CanSwitchOperationContext::attachFromOperationContext(OperationContext* opCtx) {
-    invariant(opCtx);
-    invariant(!_opCtx);
-
-    for (auto&& child : _stage->_children) {
-        child->attachFromOperationContext(opCtx);
-    }
-
-    _opCtx = opCtx;
-    doAttachFromOperationContext(opCtx);
-}
-
-void CanChangeState::saveState() {
-    _stage->_commonStats.yields++;
-    for (auto&& child : _stage->_children) {
-        child->saveState();
-    }
-
-    doSaveState();
-}
-
-void CanChangeState::restoreState() {
-    _stage->_commonStats.unyields++;
-    for (auto&& child : _stage->_children) {
-        child->restoreState();
-    }
-
-    doRestoreState();
-}
-
-void CanTrackStats::accumulate(PlanNodeId nodeId, PlanSummaryStats& summary) const {
-    if (auto stats = getSpecificStats();
-        stats && (nodeId == kEmptyPlanNodeId || _commonStats.nodeId == nodeId)) {
-        stats->accumulate(summary);
-    }
-
-    for (auto&& child : _stage->_children) {
-        child->accumulate(nodeId, summary);
-    }
-}
-}  // namespace sbe
-}  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
index 3e1712dd1e2..d818e6f6efe 100644
--- a/src/mongo/db/exec/sbe/stages/stages.h
+++ b/src/mongo/db/exec/sbe/stages/stages.h
@@ -48,12 +48,14 @@ enum class PlanState { ADVANCED, IS_EOF };
 /**
  * Provides methods to detach and re-attach to an operation context, which derived classes may
  * override to perform additional actions when these events occur.
+ *
+ * Parameter 'T' is the typename of the class derived from this class. It's used to implement the
+ * curiously recurring template pattern and access the internal state of the derived class.
  */
+template <typename T>
 class CanSwitchOperationContext {
 public:
-    CanSwitchOperationContext(PlanStage* stage) : _stage(stage) {
-        invariant(_stage);
-    }
+    CanSwitchOperationContext() = default;
 
     /**
      * Detaches from the OperationContext and releases any storage-engine state.
@@ -64,7 +66,17 @@ public:
      *
      * Propagates to all children, then calls doDetachFromOperationContext().
      */
-    void detachFromOperationContext();
+    void detachFromOperationContext() {
+        invariant(_opCtx);
+
+        auto stage = static_cast<T*>(this);
+        for (auto&& child : stage->_children) {
+            child->detachFromOperationContext();
+        }
+
+        stage->doDetachFromOperationContext();
+        _opCtx = nullptr;
+    }
 
     /**
      * Reattaches to the OperationContext and reacquires any storage-engine state.
@@ -74,29 +86,35 @@ public:
      *
      * Propagates to all children, then calls doReattachToOperationContext().
      */
-    void attachFromOperationContext(OperationContext* opCtx);
+    void attachFromOperationContext(OperationContext* opCtx) {
+        invariant(opCtx);
+        invariant(!_opCtx);
 
-protected:
-    // Derived classes can optionally override these methods.
-    virtual void doDetachFromOperationContext() {}
-    virtual void doAttachFromOperationContext(OperationContext* opCtx) {}
+        auto stage = static_cast<T*>(this);
+        for (auto&& child : stage->_children) {
+            child->attachFromOperationContext(opCtx);
+        }
 
-    OperationContext* _opCtx{nullptr};
+        _opCtx = opCtx;
+        stage->doAttachFromOperationContext(opCtx);
+    }
 
-private:
-    PlanStage* const _stage;
+protected:
+    OperationContext* _opCtx{nullptr};
 };
 
 /**
  * Provides methods to save and restore the state of the object which derives from this class
  * when corresponding events are generated as a response to a change in the underlying data source.
  * Derived classes may override these methods to perform additional actions when these events occur.
+ *
+ * Parameter 'T' is the typename of the class derived from this class. It's used to implement the
+ * curiously recurring template pattern and access the internal state of the derived class.
  */
+template <typename T>
 class CanChangeState {
 public:
-    CanChangeState(PlanStage* stage) : _stage(stage) {
-        invariant(_stage);
-    }
+    CanChangeState() = default;
 
     /**
      * Notifies the stage that the underlying data source may change.
@@ -106,7 +124,15 @@ public:
      *
      * Propagates to all children, then calls doSaveState().
      */
-    void saveState();
+    void saveState() {
+        auto stage = static_cast<T*>(this);
+        stage->_commonStats.yields++;
+        for (auto&& child : stage->_children) {
+            child->saveState();
+        }
+
+        stage->doSaveState();
+    }
 
     /**
      * Notifies the stage that underlying data is stable again and prepares for calls to work().
@@ -119,24 +145,27 @@ public:
      * collection drop. May throw a WriteConflictException, in which case the caller may choose to
      * retry.
      */
-    void restoreState();
-
-protected:
-    // Derived classes can optionally override these methods.
-    virtual void doSaveState() {}
-    virtual void doRestoreState() {}
+    void restoreState() {
+        auto stage = static_cast<T*>(this);
+        stage->_commonStats.unyields++;
+        for (auto&& child : stage->_children) {
+            child->restoreState();
+        }
 
-private:
-    PlanStage* const _stage;
+        stage->doRestoreState();
+    }
 };
 
 /**
  * Provides methods to obtain execution statistics specific to a plan stage.
+ *
+ * Parameter 'T' is the typename of the class derived from this class. It's used to implement the
+ * curiously recurring template pattern and access the internal state of the derived class.
  */
+template <typename T>
 class CanTrackStats {
 public:
-    CanTrackStats(PlanStage* stage, StringData stageType, PlanNodeId nodeId)
-        : _stage(stage), _commonStats(stageType, nodeId) {}
+    CanTrackStats(StringData stageType, PlanNodeId nodeId) : _commonStats(stageType, nodeId) {}
 
     /**
      * Returns a tree of stats. If the stage has any children it must propagate the request for
@@ -171,7 +200,17 @@ public:
      * invoking 'accumulate(summary)' on the SpecificStats instance obtained by calling
      * 'getSpecificStats()'.
      */
-    void accumulate(PlanNodeId nodeId, PlanSummaryStats& summary) const;
+    void accumulate(PlanNodeId nodeId, PlanSummaryStats& summary) const {
+        if (auto stats = getSpecificStats();
+            stats && (nodeId == kEmptyPlanNodeId || _commonStats.nodeId == nodeId)) {
+            stats->accumulate(summary);
+        }
+
+        auto stage = static_cast<const T*>(this);
+        for (auto&& child : stage->_children) {
+            child->accumulate(nodeId, summary);
+        }
+    }
 
 protected:
     PlanState trackPlanState(PlanState state) {
@@ -184,7 +223,6 @@ protected:
         return state;
     }
 
-    PlanStage* const _stage;
     CommonStats _commonStats;
 };
 
@@ -229,16 +267,13 @@ private:
 /**
  * This is an abstract base class of all plan stages in SBE.
  */
-class PlanStage : public CanSwitchOperationContext,
-                  public CanChangeState,
-                  public CanTrackStats,
+class PlanStage : public CanSwitchOperationContext<PlanStage>,
+                  public CanChangeState<PlanStage>,
+                  public CanTrackStats<PlanStage>,
                   public CanInterrupt {
 public:
     PlanStage(StringData stageType, PlanYieldPolicy* yieldPolicy, PlanNodeId nodeId)
-        : CanSwitchOperationContext{this},
-          CanChangeState{this},
-          CanTrackStats{this, stageType, nodeId},
-          CanInterrupt{yieldPolicy} {}
+        : CanTrackStats{stageType, nodeId}, CanInterrupt{yieldPolicy} {}
 
     PlanStage(StringData stageType, PlanNodeId nodeId) : PlanStage(stageType, nullptr, nodeId) {}
 
@@ -290,11 +325,17 @@ public:
         return {DebugPrinter::Block(str)};
     }
 
-    friend class CanSwitchOperationContext;
-    friend class CanChangeState;
-    friend class CanTrackStats;
+    friend class CanSwitchOperationContext<PlanStage>;
+    friend class CanChangeState<PlanStage>;
+    friend class CanTrackStats<PlanStage>;
 
 protected:
+    // Derived classes can optionally override these methods.
+    virtual void doSaveState() {}
+    virtual void doRestoreState() {}
+    virtual void doDetachFromOperationContext() {}
+    virtual void doAttachFromOperationContext(OperationContext* opCtx) {}
+
     std::vector<std::unique_ptr<PlanStage>> _children;
 };