SERVER-70490 Update cost model on Mongod startup

4 files changed, 73 insertions, 2 deletions

diff --git a/jstests/noPassthrough/start_up_with_custom_cost_model.js b/jstests/noPassthrough/start_up_with_custom_cost_model.js
new file mode 100644
index 00000000000..85d1098cd86
--- /dev/null
+++ b/jstests/noPassthrough/start_up_with_custom_cost_model.js
@@ -0,0 +1,53 @@
+/**
+ * Tests that 'internalCostModelCoefficients' can be set on startup.
+ */
+(function() {
+'use strict';
+
+load("jstests/libs/optimizer_utils.js");  // For checkCascadesOptimizerEnabled.
+
+function getScanCostWith(customScanCost) {
+    const costStr = `{"scanIncrementalCost": ${customScanCost}}`;
+    const conn = MongoRunner.runMongod({setParameter: {'internalCostModelCoefficients': costStr}});
+
+    const db = conn.getDB(jsTestName());
+    const coll = db.start_up_with_custom_cost_model;
+
+    if (!checkCascadesOptimizerEnabled(db)) {
+        jsTestLog("Skipping test because the optimizer is not enabled");
+        MongoRunner.stopMongod(conn);
+        return;
+    }
+
+    const nDocuments = 100;
+    assert.commandWorked(coll.insert(Array.from({length: nDocuments}, (_, i) => {
+        return {a: 3, b: 3, c: i};
+    })));
+
+    // Get scan cost.
+    const explain = coll.explain("executionStats").aggregate([]);
+    assert.eq(nDocuments, explain.executionStats.nReturned);
+    const internalCost =
+        assert.commandWorked(db.adminCommand({getParameter: 1, internalCostModelCoefficients: 1}));
+
+    assert(internalCost.hasOwnProperty("internalCostModelCoefficients"));
+    assert.eq(internalCost['internalCostModelCoefficients'], costStr);
+
+    const scanNode = navigateToPlanPath(explain, "child");
+    assertValueOnPath("PhysicalScan", scanNode, "nodeType");
+
+    const scanCost = scanNode.properties.cost;
+
+    MongoRunner.stopMongod(conn);
+
+    return scanCost;
+}
+
+const scanCost1 = getScanCostWith(0.2);
+const scanCost2 = getScanCostWith(0.4);
+if (scanCost1 === undefined) {
+    return;
+}
+
+assert.lt(scanCost1, scanCost2);
+}());
diff --git a/src/mongo/db/query/cost_model/cost_model_on_update.cpp b/src/mongo/db/query/cost_model/cost_model_on_update.cpp
index 290b3bc706d..c57d00b4140 100644
--- a/src/mongo/db/query/cost_model/cost_model_on_update.cpp
+++ b/src/mongo/db/query/cost_model/cost_model_on_update.cpp
@@ -33,6 +33,9 @@
 #include "mongo/db/client.h"
 #include "mongo/db/query/cost_model/cost_model_manager.h"
 #include "mongo/db/query/query_knobs_gen.h"
+#include "mongo/logv2/log.h"
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery
 
 namespace mongo::cost_model {
 
@@ -55,7 +58,8 @@ Status updateCostCoefficients() {
         auto updater = onCoefficientsChangeUpdater(serviceCtx).get();
         updater->updateCoefficients(serviceCtx, overrides);
     } else {
-        tasserted(7049001, "Client must be non null");
+        // 'client' may be null if the server parameter is set on mongod startup.
+        LOGV2_DEBUG(7049001, 5, "Cost model coefficients updated on startup");
     }
 
     return Status::OK();
diff --git a/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.cpp b/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.cpp
index 7ef16963e7b..aa8b430fc78 100644
--- a/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.cpp
+++ b/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.cpp
@@ -29,6 +29,9 @@
 
 #include "mongo/db/query/cost_model/on_coefficients_change_updater_impl.h"
 
+#include "mongo/bson/json.h"
+#include "mongo/db/query/query_knobs_gen.h"
+
 namespace mongo::cost_model {
 
 const Decorable<ServiceContext>::Decoration<CostModelManager> costModelManager =
@@ -41,7 +44,16 @@ void OnCoefficientsChangeUpdaterImpl::updateCoefficients(ServiceContext* service
 
 ServiceContext::ConstructorActionRegisterer costModelUpdaterRegisterer{
     "costModelUpdaterRegisterer", [](ServiceContext* serviceCtx) {
+        BSONObj overrides = internalCostModelCoefficients.empty()
+            ? BSONObj()
+            : fromjson(internalCostModelCoefficients);
+
         onCoefficientsChangeUpdater(serviceCtx) =
-            std::make_unique<OnCoefficientsChangeUpdaterImpl>();
+            std::make_unique<OnCoefficientsChangeUpdaterImpl>(serviceCtx, overrides);
     }};
+
+OnCoefficientsChangeUpdaterImpl::OnCoefficientsChangeUpdaterImpl(ServiceContext* serviceCtx,
+                                                                 const BSONObj& overrides) {
+    updateCoefficients(serviceCtx, overrides);
+}
 }  // namespace mongo::cost_model
diff --git a/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.h b/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.h
index b19b5f7d333..f49a5c73bc7 100644
--- a/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.h
+++ b/src/mongo/db/query/cost_model/on_coefficients_change_updater_impl.h
@@ -38,6 +38,8 @@ extern const Decorable<ServiceContext>::Decoration<CostModelManager> costModelMa
 
 class OnCoefficientsChangeUpdaterImpl final : public OnCoefficientsChangeUpdater {
 public:
+    OnCoefficientsChangeUpdaterImpl(ServiceContext* serviceCtx, const BSONObj& overrides);
+
     void updateCoefficients(ServiceContext* serviceCtx, const BSONObj& overrides) final;
 };
 }  // namespace mongo::cost_model