mongo: SERVER-71230 Add baton support to async rpc senders

Commit: 99cf1c7c97eb67ad7505140e20cab9d1c23d0b0c

Author: George Wangensteen

mongo/src/mongo/executor/async_rpc.cpp

@@ -53,18 +53,25 @@ class AsyncRPCRunnerImpl : public AsyncRPCRunner {
                                                           Targeter* targeter,
                                                           OperationContext* opCtx,
                                                           std::shared_ptr<TaskExecutor> exec,
-                                                          CancellationToken token) final {
+                                                          CancellationToken token,
+                                                          BatonHandle baton) final {
+        auto proxyExec = std::make_shared<ProxyingExecutor>(baton, exec);
         auto targetsUsed = std::make_shared<std::vector<HostAndPort>>();
         return targeter->resolve(token)
-            .thenRunOn(exec)
-            .then([dbName, cmdBSON, opCtx, exec = std::move(exec), token, targetsUsed](
-                      std::vector<HostAndPort> targets) {
+            .thenRunOn(proxyExec)
+            .then([dbName,
+                   cmdBSON,
+                   opCtx,
+                   exec = std::move(exec),
+                   token,
+                   baton = std::move(baton),
+                   targetsUsed](std::vector<HostAndPort> targets) {
                 invariant(targets.size(),
                           "Successful targeting implies there are hosts to target.");
                 *targetsUsed = targets;
                 executor::RemoteCommandRequestOnAny executorRequest(
                     targets, dbName.toString(), cmdBSON, rpc::makeEmptyMetadata(), opCtx);
-                return exec->scheduleRemoteCommandOnAny(executorRequest, token);
+                return exec->scheduleRemoteCommandOnAny(executorRequest, token, std::move(baton));
             })
             .onError([targetsUsed](Status s) -> StatusWith<TaskExecutor::ResponseOnAnyStatus> {
                 // If there was a scheduling error or other local error before the

mongo/src/mongo/executor/async_rpc.h

@@ -86,12 +86,14 @@ struct AsyncRPCOptions {
     AsyncRPCOptions(CommandType cmd,
                     std::shared_ptr<executor::TaskExecutor> exec,
                     CancellationToken token,
-                    std::shared_ptr<RetryPolicy> retryPolicy = std::make_shared<NeverRetryPolicy>())
-        : cmd{cmd}, exec{exec}, token{token}, retryPolicy{retryPolicy} {}
+                    std::shared_ptr<RetryPolicy> retryPolicy = std::make_shared<NeverRetryPolicy>(),
+                    BatonHandle baton = nullptr)
+        : cmd{cmd}, exec{exec}, token{token}, retryPolicy{retryPolicy}, baton{std::move(baton)} {}
     CommandType cmd;
     std::shared_ptr<executor::TaskExecutor> exec;
     CancellationToken token;
     std::shared_ptr<RetryPolicy> retryPolicy;
+    BatonHandle baton;
 };
 
 /**
@@ -124,7 +126,22 @@ class AsyncRPCRunner {
         Targeter* targeter,
         OperationContext* opCtx,
         std::shared_ptr<TaskExecutor> exec,
-        CancellationToken token) = 0;
+        CancellationToken token,
+        BatonHandle baton) = 0;
+    ExecutorFuture<AsyncRPCInternalResponse> _sendCommand(StringData dbName,
+                                                          BSONObj cmdBSON,
+                                                          Targeter* targeter,
+                                                          OperationContext* opCtx,
+                                                          std::shared_ptr<TaskExecutor> exec,
+                                                          CancellationToken token) {
+        return _sendCommand(std::move(dbName),
+                            std::move(cmdBSON),
+                            std::move(targeter),
+                            std::move(opCtx),
+                            std::move(exec),
+                            std::move(token),
+                            nullptr);
+    }
     static AsyncRPCRunner* get(ServiceContext* serviceContext);
     static void set(ServiceContext* serviceContext, std::unique_ptr<AsyncRPCRunner> theRunner);
 };
@@ -154,6 +171,31 @@ struct RetryDelayAsBackoff {
     RetryPolicy* _policy;
 };
 
+class ProxyingExecutor : public OutOfLineExecutor,
+                         public std::enable_shared_from_this<ProxyingExecutor> {
+public:
+    ProxyingExecutor(BatonHandle baton, std::shared_ptr<TaskExecutor> executor)
+        : _baton{std::move(baton)}, _executor{std::move(executor)} {}
+
+    void schedule(Task func) override {
+        if (_baton)
+            return _baton->schedule(std::move(func));
+        return _executor->schedule(std::move(func));
+    }
+
+    ExecutorFuture<void> sleepFor(Milliseconds duration, const CancellationToken& token) {
+        auto deadline = Date_t::now() + duration;
+        if (auto netBaton = _baton ? _baton->networking() : nullptr; netBaton) {
+            return netBaton->waitUntil(deadline, token).thenRunOn(shared_from_this());
+        }
+        return _executor->sleepFor(duration, token);
+    }
+
+private:
+    BatonHandle _baton;
+    std::shared_ptr<TaskExecutor> _executor;
+};
+
 template <typename CommandType>
 ExecutorFuture<AsyncRPCResponse<typename CommandType::Reply>> sendCommandWithRunner(
     BSONObj cmdBSON,
@@ -162,6 +204,7 @@ ExecutorFuture<AsyncRPCResponse<typename CommandType::Reply>> sendCommandWithRun
     OperationContext* opCtx,
     std::unique_ptr<Targeter> targeter) {
     using ReplyType = AsyncRPCResponse<typename CommandType::Reply>;
+    auto proxyExec = std::make_shared<ProxyingExecutor>(options->baton, options->exec);
     auto tryBody = [=, targeter = std::move(targeter)] {
         // Execute the command after extracting the db name and bson from the CommandType.
         // Wrapping this function allows us to separate the CommandType parsing logic from the
@@ -181,7 +224,8 @@ ExecutorFuture<AsyncRPCResponse<typename CommandType::Reply>> sendCommandWithRun
                 return shouldStopRetry;
             })
             .withBackoffBetweenIterations(RetryDelayAsBackoff(options->retryPolicy.get()))
-            .on(options->exec, CancellationToken::uncancelable());
+            .on(proxyExec, CancellationToken::uncancelable());
+
     return std::move(resFuture)
         .then([](detail::AsyncRPCInternalResponse r) -> ReplyType {
             auto res = CommandType::Reply::parseSharingOwnership(IDLParserContext("AsyncRPCRunner"),

mongo/src/mongo/executor/async_rpc_test.cpp

@@ -44,6 +44,7 @@
 #include "mongo/executor/thread_pool_task_executor_test_fixture.h"
 #include "mongo/rpc/topology_version_gen.h"
 #include "mongo/unittest/bson_test_util.h"
+#include "mongo/unittest/thread_assertion_monitor.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/duration.h"
 #include "mongo/util/future.h"
@@ -412,6 +413,44 @@ TEST_F(AsyncRPCTestFixture, ExecutorShutdown) {
     ASSERT(ErrorCodes::isA<ErrorCategory::CancellationError>(extraInfo->asLocal()));
 }
 
+TEST_F(AsyncRPCTestFixture, BatonTest) {
+    std::unique_ptr<Targeter> targeter = std::make_unique<LocalHostTargeter>();
+    auto retryPolicy = std::make_shared<NeverRetryPolicy>();
+    HelloCommand helloCmd;
+    HelloCommandReply helloReply = HelloCommandReply(TopologyVersion(OID::gen(), 0));
+    initializeCommand(helloCmd);
+    auto opCtxHolder = makeOperationContext();
+    auto baton = opCtxHolder->getBaton();
+    auto options = std::make_shared<AsyncRPCOptions<HelloCommand>>(
+        helloCmd, getExecutorPtr(), _cancellationToken);
+    options->baton = baton;
+    auto resultFuture = sendCommand(options, opCtxHolder.get(), std::move(targeter));
+
+    Notification<void> seenNetworkRequest;
+    unittest::ThreadAssertionMonitor monitor;
+    // This thread will respond to the request we sent via sendCommand above.
+    auto networkResponder = monitor.spawn([&] {
+        onCommand([&](const auto& request) {
+            ASSERT(request.cmdObj["hello"]);
+            seenNetworkRequest.set();
+            monitor.notifyDone();
+            return helloReply.toBSON();
+        });
+    });
+    // Wait on the opCtx until networkResponder has observed the network request.
+    // While we block on the opCtx, the current thread should run jobs scheduled
+    // on the baton, including enqueuing the network request via `sendCommand` above.
+    seenNetworkRequest.get(opCtxHolder.get());
+
+    networkResponder.join();
+    // Wait on the opCtx again to allow the current thread, via the baton, to propogate
+    // the network response up into the resultFuture.
+    AsyncRPCResponse res = resultFuture.get(opCtxHolder.get());
+
+    ASSERT_BSONOBJ_EQ(res.response.toBSON(), helloReply.toBSON());
+    ASSERT_EQ(HostAndPort("localhost", serverGlobalParams.port), res.targetUsed);
+}
+
 /*
  * Basic Targeter that returns the host that invoked it.
  */
@@ -507,6 +546,42 @@ TEST_F(AsyncRPCTestFixture, FailedTargeting) {
     ASSERT(extraInfo->isLocal());
     ASSERT_EQ(extraInfo->asLocal(), targeterFailStatus);
 }
+TEST_F(AsyncRPCTestFixture, BatonShutdownExecutorAlive) {
+    std::unique_ptr<Targeter> targeter = std::make_unique<LocalHostTargeter>();
+    auto retryPolicy = std::make_shared<TestRetryPolicy>();
+    const auto maxNumRetries = 5;
+    const auto retryDelay = Milliseconds(10);
+    retryPolicy->setMaxNumRetries(maxNumRetries);
+    for (int i = 0; i < maxNumRetries; ++i)
+        retryPolicy->pushRetryDelay(retryDelay);
+    HelloCommand helloCmd;
+    HelloCommandReply helloReply = HelloCommandReply(TopologyVersion(OID::gen(), 0));
+    initializeCommand(helloCmd);
+    auto opCtxHolder = makeOperationContext();
+    auto subBaton = opCtxHolder->getBaton()->makeSubBaton();
+    auto options = std::make_shared<AsyncRPCOptions<HelloCommand>>(
+        helloCmd, getExecutorPtr(), _cancellationToken);
+    options->baton = *subBaton;
+    auto resultFuture = sendCommand(options, opCtxHolder.get(), std::move(targeter));
+
+    subBaton.shutdown();
+
+    auto error = resultFuture.getNoThrow().getStatus();
+    auto expectedDetachError = Status(ErrorCodes::ShutdownInProgress, "Baton detached");
+    auto expectedOuterReason = "Remote command execution failed due to executor shutdown";
+
+    ASSERT_EQ(error.code(), ErrorCodes::RemoteCommandExecutionError);
+    ASSERT_EQ(error.reason(), expectedOuterReason);
+
+    auto extraInfo = error.extraInfo<AsyncRPCErrorInfo>();
+    ASSERT(extraInfo);
+
+    ASSERT(extraInfo->isLocal());
+    auto localError = extraInfo->asLocal();
+    ASSERT_EQ(localError, expectedDetachError);
+
+    ASSERT_EQ(0, retryPolicy->getNumRetriesPerformed());
+}
 
 TEST_F(AsyncRPCTestFixture, SendTxnCommandWithoutTxnRouterAppendsNoTxnFields) {
     ShardId shardId("shard");

mongo/src/mongo/executor/hedged_async_rpc.h

@@ -115,22 +115,23 @@ SemiFuture<AsyncRPCResponse<typename CommandType::Reply>> sendHedgedCommand(
     std::shared_ptr<executor::TaskExecutor> exec,
     CancellationToken token,
     std::shared_ptr<RetryPolicy> retryPolicy = std::make_shared<NeverRetryPolicy>(),
-    ReadPreferenceSetting readPref = ReadPreferenceSetting(ReadPreference::PrimaryOnly)) {
+    ReadPreferenceSetting readPref = ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+    BatonHandle baton = nullptr) {
     using SingleResponse = AsyncRPCResponse<typename CommandType::Reply>;
 
     // Set up cancellation token to cancel remaining hedged operations.
     CancellationSource hedgeCancellationToken{token};
     auto targetsAttempted = std::make_shared<std::vector<HostAndPort>>();
+    auto proxyExec = std::make_shared<detail::ProxyingExecutor>(baton, exec);
     auto tryBody = [=, targeter = std::move(targeter)] {
         return targeter->resolve(token)
-            .thenRunOn(exec)
+            .thenRunOn(proxyExec)
             .onError([](Status status) -> StatusWith<std::vector<HostAndPort>> {
                 // Targeting error; rewrite it to a RemoteCommandExecutionError and skip
                 // command execution body. We'll retry if the policy indicates to.
                 return Status{AsyncRPCErrorInfo(status), status.reason()};
             })
-            .then([cmd, opCtx, exec, token, hedgeCancellationToken, readPref, targetsAttempted](
-                      std::vector<HostAndPort> targets) {
+            .then([=](std::vector<HostAndPort> targets) {
                 invariant(targets.size(),
                           "Successful targeting implies there are hosts to target.");
                 *targetsAttempted = targets;
@@ -146,8 +147,9 @@ SemiFuture<AsyncRPCResponse<typename CommandType::Reply>> sendHedgedCommand(
                     std::unique_ptr<Targeter> t = std::make_unique<FixedTargeter>(targets[i]);
                     auto options = std::make_shared<AsyncRPCOptions<CommandType>>(
                         cmd, exec, hedgeCancellationToken.token());
-                    requests.emplace_back(
-                        sendCommand(options, opCtx, std::move(t)).thenRunOn(exec));
+                    options->baton = baton;
+                    requests.push_back(
+                        sendCommand(options, opCtx, std::move(t)).thenRunOn(proxyExec));
                 }
 
                 /**
@@ -189,7 +191,7 @@ SemiFuture<AsyncRPCResponse<typename CommandType::Reply>> sendHedgedCommand(
                 !retryPolicy->recordAndEvaluateRetry(swResponse.getStatus());
         })
         .withBackoffBetweenIterations(detail::RetryDelayAsBackoff(retryPolicy.get()))
-        .on(exec, CancellationToken::uncancelable())
+        .on(proxyExec, CancellationToken::uncancelable())
         // We go inline here to intercept executor-shutdown errors and re-write them
         // so that the API always returns RemoteCommandExecutionError. Additionally,
         // we need to make sure we cancel outstanding requests.