SERVER-49107 Add support for async execution to MongoD command path

1 files changed, 50 insertions, 0 deletions

diff --git a/src/mongo/db/commands.h b/src/mongo/db/commands.h
index d1cf2b369c8..e1805991cab 100644
--- a/src/mongo/db/commands.h
+++ b/src/mongo/db/commands.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include <boost/optional.hpp>
+#include <fmt/format.h>
 #include <functional>
 #include <string>
 #include <vector>
@@ -47,10 +48,12 @@
 #include "mongo/db/query/explain.h"
 #include "mongo/db/read_concern_support_result.h"
 #include "mongo/db/repl/read_concern_args.h"
+#include "mongo/db/request_execution_context.h"
 #include "mongo/db/write_concern.h"
 #include "mongo/rpc/op_msg.h"
 #include "mongo/rpc/reply_builder_interface.h"
 #include "mongo/util/fail_point.h"
+#include "mongo/util/future.h"
 #include "mongo/util/string_map.h"
 
 namespace mongo {
@@ -92,11 +95,27 @@ public:
                              CommandInvocation* invocation) = 0;
 
     /**
+     * A behavior to perform before CommandInvocation::asyncRun(). Defaults to `onBeforeRun(...)`.
+     */
+    virtual void onBeforeAsyncRun(std::shared_ptr<RequestExecutionContext> rec,
+                                  CommandInvocation* invocation) {
+        onBeforeRun(rec->getOpCtx(), rec->getRequest(), invocation);
+    }
+
+    /**
      * A behavior to perform after CommandInvocation::run()
      */
     virtual void onAfterRun(OperationContext* opCtx,
                             const OpMsgRequest& request,
                             CommandInvocation* invocation) = 0;
+
+    /**
+     * A behavior to perform after CommandInvocation::asyncRun(). Defaults to `onAfterRun(...)`.
+     */
+    virtual void onAfterAsyncRun(std::shared_ptr<RequestExecutionContext> rec,
+                                 CommandInvocation* invocation) {
+        onAfterRun(rec->getOpCtx(), rec->getRequest(), invocation);
+    }
 };
 
 // Various helpers unrelated to any single command or to the command registry.
@@ -236,6 +255,15 @@ struct CommandHelpers {
                                      rpc::ReplyBuilderInterface* response);
 
     /**
+     * Runs a previously parsed command and propagates the result to the ReplyBuilderInterface. For
+     * commands that do not offer an implementation tailored for asynchronous execution, the future
+     * schedules the execution of the default implementation, historically designed for synchronous
+     * execution.
+     */
+    static Future<void> runCommandInvocationAsync(std::shared_ptr<RequestExecutionContext> rec,
+                                                  std::shared_ptr<CommandInvocation> invocation);
+
+    /**
      * If '!invocation', we're logging about a Command pre-parse. It has to punt on the logged
      * namespace, giving only the request's $db. Since the Command hasn't parsed the request body,
      * we can't know the collection part of that namespace, so we leave it blank in the audit log.
@@ -568,6 +596,16 @@ public:
      */
     virtual void run(OperationContext* opCtx, rpc::ReplyBuilderInterface* result) = 0;
 
+    /**
+     * Returns a future that can schedule asynchronous execution of the command. By default, the
+     * future falls back to the execution of `run(...)`, thus the default semantics of
+     * `runAsync(...)` is identical to that of `run(...).
+     */
+    virtual Future<void> runAsync(std::shared_ptr<RequestExecutionContext> rec) {
+        run(rec->getOpCtx(), rec->getReplyBuilder());
+        return Status::OK();
+    }
+
     virtual void explain(OperationContext* opCtx,
                          ExplainOptions::Verbosity verbosity,
                          rpc::ReplyBuilderInterface* result) {
@@ -705,6 +743,18 @@ public:
                                      rpc::ReplyBuilderInterface* replyBuilder) = 0;
 
     /**
+     * Provides a future that may run the command asynchronously. By default, it falls back to
+     * runWithReplyBuilder.
+     */
+    virtual Future<void> runAsync(std::shared_ptr<RequestExecutionContext> rec, std::string db) {
+        if (!runWithReplyBuilder(
+                rec->getOpCtx(), db, rec->getRequest().body, rec->getReplyBuilder()))
+            return Status(ErrorCodes::FailedToRunWithReplyBuilder,
+                          fmt::format("Failed to run command: {}", rec->getCommand()->getName()));
+        return Status::OK();
+    }
+
+    /**
      * Commands which can be explained override this method. Any operation which has a query
      * part and executes as a tree of execution stages can be explained. A command should
      * implement explain by: