SERVER-35708 Make mongos attach errorLabels on command response

7 files changed, 132 insertions, 33 deletions

diff --git a/buildscripts/resmokeconfig/suites/sharded_core_txns.yml b/buildscripts/resmokeconfig/suites/sharded_core_txns.yml
index 44d9f5774ed..42756519afd 100644
--- a/buildscripts/resmokeconfig/suites/sharded_core_txns.yml
+++ b/buildscripts/resmokeconfig/suites/sharded_core_txns.yml
@@ -53,7 +53,7 @@ selector:
   # TODO SERVER-33709: Support readConcern snapshot in cluster count command.
   - jstests/core/txns/commands_not_allowed_in_txn.js
 
-  # TODO SERVER-35708 Mongos does not attach errorLabels on command response.
+  # TODO SERVER-36580 Mongos should preserve error labels for write ops.
   - jstests/core/txns/multi_statement_transaction_abort.js
 
   # TODO SERVER-35825: Mongos should only allow readConcern on writes that start transactions.
diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 83d11b5cf3b..a8d843ba7d0 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -799,6 +799,7 @@ env.Library(
         '$BUILD_DIR/mongo/base',
     ],
     LIBDEPS_PRIVATE=[
+        'handle_request_response',
         'snapshot_window_util',
         '$BUILD_DIR/mongo/db/auth/auth',
         '$BUILD_DIR/mongo/db/auth/authprivilege',
@@ -1383,6 +1384,16 @@ env.Library(
 )
 
 env.Library(
+    target='handle_request_response',
+    source=[
+        'handle_request_response.cpp',
+    ],
+    LIBDEPS=[
+        'logical_session_cache_impl',
+    ],
+)
+
+env.Library(
     target='transaction_reaper',
     source=[
         'transaction_reaper.cpp',
diff --git a/src/mongo/db/handle_request_response.cpp b/src/mongo/db/handle_request_response.cpp
new file mode 100644
index 00000000000..877baef9c77
--- /dev/null
+++ b/src/mongo/db/handle_request_response.cpp
@@ -0,0 +1,57 @@
+/**
+ * Copyright (C) 2017 MongoDB Inc.
+ *
+ * This program is free software: you can redistribute it and/or  modify
+ * it under the terms of the GNU Affero General Public License, version 3,
+ * as published by the Free Software Foundation.
+ *
+ * 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
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ * 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 GNU Affero General 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/db/handle_request_response.h"
+
+namespace mongo {
+
+BSONObj getErrorLabels(const boost::optional<OperationSessionInfoFromClient>& sessionOptions,
+                       const std::string& commandName,
+                       ErrorCodes::Error code) {
+    // By specifying "autocommit", the user indicates they want to run a transaction.
+    if (!sessionOptions || !sessionOptions->getAutocommit()) {
+        return {};
+    }
+
+    bool isRetryable = ErrorCodes::isNotMasterError(code) || ErrorCodes::isShutdownError(code);
+    bool isTransientTransactionError = code == ErrorCodes::WriteConflict  //
+        || code == ErrorCodes::SnapshotUnavailable                        //
+        || code == ErrorCodes::NoSuchTransaction                          //
+        || code == ErrorCodes::LockTimeout                                //
+        || code == ErrorCodes::PreparedTransactionInProgress              //
+        // Clients can retry a single commitTransaction command, but cannot retry the whole
+        // transaction if commitTransaction fails due to NotMaster.
+        || (isRetryable && (commandName != "commitTransaction"));
+
+    if (isTransientTransactionError) {
+        return BSON("errorLabels" << BSON_ARRAY("TransientTransactionError"));
+    }
+    return {};
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/handle_request_response.h b/src/mongo/db/handle_request_response.h
new file mode 100644
index 00000000000..e561dc02fd6
--- /dev/null
+++ b/src/mongo/db/handle_request_response.h
@@ -0,0 +1,42 @@
+/**
+ * Copyright (C) 2017 MongoDB Inc.
+ *
+ * This program is free software: you can redistribute it and/or  modify
+ * it under the terms of the GNU Affero General Public License, version 3,
+ * as published by the Free Software Foundation.
+ *
+ * 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
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ * 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 GNU Affero General 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.
+ */
+
+#pragma once
+
+#include "mongo/db/logical_session_id.h"
+
+namespace mongo {
+
+/**
+ * Returns the error labels for the given error.
+ */
+BSONObj getErrorLabels(const boost::optional<OperationSessionInfoFromClient>& sessionOptions,
+                       const std::string& commandName,
+                       ErrorCodes::Error code);
+
+}  // namespace mongo
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index b8b8e3b68f5..fd185ca00fe 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/curop_metrics.h"
 #include "mongo/db/cursor_manager.h"
 #include "mongo/db/dbdirectclient.h"
+#include "mongo/db/handle_request_response.h"
 #include "mongo/db/initialize_operation_session_info.h"
 #include "mongo/db/introspect.h"
 #include "mongo/db/jsobj.h"
@@ -220,30 +221,6 @@ void generateErrorResponse(OperationContext* opCtx,
     replyBuilder->getBodyBuilder().appendElements(replyMetadata);
 }
 
-BSONObj getErrorLabels(const boost::optional<OperationSessionInfoFromClient>& sessionOptions,
-                       const std::string& commandName,
-                       ErrorCodes::Error code) {
-    // By specifying "autocommit", the user indicates they want to run a transaction.
-    if (!sessionOptions || !sessionOptions->getAutocommit()) {
-        return {};
-    }
-
-    bool isRetryable = ErrorCodes::isNotMasterError(code) || ErrorCodes::isShutdownError(code);
-    bool isTransientTransactionError = code == ErrorCodes::WriteConflict  //
-        || code == ErrorCodes::SnapshotUnavailable                        //
-        || code == ErrorCodes::NoSuchTransaction                          //
-        || code == ErrorCodes::LockTimeout                                //
-        || code == ErrorCodes::PreparedTransactionInProgress              //
-        // Clients can retry a single commitTransaction command, but cannot retry the whole
-        // transaction if commitTransaction fails due to NotMaster.
-        || (isRetryable && (commandName != "commitTransaction"));
-
-    if (isTransientTransactionError) {
-        return BSON("errorLabels" << BSON_ARRAY("TransientTransactionError"));
-    }
-    return {};
-}
-
 /**
  * Guard object for making a good-faith effort to enter maintenance mode and leave it when it
  * goes out of scope.
diff --git a/src/mongo/s/commands/SConscript b/src/mongo/s/commands/SConscript
index 2c642d077b9..5c75e96236e 100644
--- a/src/mongo/s/commands/SConscript
+++ b/src/mongo/s/commands/SConscript
@@ -115,6 +115,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/commands/test_commands_enabled',
         '$BUILD_DIR/mongo/db/commands/write_commands_common',
         '$BUILD_DIR/mongo/db/ftdc/ftdc_server',
+        '$BUILD_DIR/mongo/db/handle_request_response',
         '$BUILD_DIR/mongo/db/logical_session_cache_impl',
         '$BUILD_DIR/mongo/db/pipeline/aggregation',
         '$BUILD_DIR/mongo/db/pipeline/mongos_process_interface',
diff --git a/src/mongo/s/commands/strategy.cpp b/src/mongo/s/commands/strategy.cpp
index e10777fcde2..efa76197bb3 100644
--- a/src/mongo/s/commands/strategy.cpp
+++ b/src/mongo/s/commands/strategy.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/commands.h"
 #include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/curop.h"
+#include "mongo/db/handle_request_response.h"
 #include "mongo/db/initialize_operation_session_info.h"
 #include "mongo/db/lasterror.h"
 #include "mongo/db/logical_clock.h"
@@ -287,10 +288,15 @@ void execCommandClient(OperationContext* opCtx,
 
 MONGO_FAIL_POINT_DEFINE(doNotRefreshShardsOnRetargettingError);
 
+/**
+ * Executes the command for the given request, and appends the result to replyBuilder
+ * and error labels, if any, to errorBuilder.
+ */
 void runCommand(OperationContext* opCtx,
                 const OpMsgRequest& request,
                 const NetworkOp opType,
-                rpc::ReplyBuilderInterface* replyBuilder) {
+                rpc::ReplyBuilderInterface* replyBuilder,
+                BSONObjBuilder* errorBuilder) {
     auto const commandName = request.getCommandName();
     auto const command = CommandHelpers::findCommand(commandName);
     if (!command) {
@@ -343,10 +349,11 @@ void runCommand(OperationContext* opCtx,
     }
 
     boost::optional<ScopedRouterSession> scopedSession;
-    if (auto osi = initializeOperationSessionInfo(
-            opCtx, request.body, command->requiresAuth(), true, true, true)) {
+    auto osi = initializeOperationSessionInfo(
+        opCtx, request.body, command->requiresAuth(), true, true, true);
 
-        if (osi->getAutocommit()) {
+    try {
+        if (osi && osi->getAutocommit()) {
             scopedSession.emplace(opCtx);
 
             auto txnRouter = TransactionRouter::get(opCtx);
@@ -360,9 +367,7 @@ void runCommand(OperationContext* opCtx,
 
             txnRouter->beginOrContinueTxn(opCtx, *txnNumber, startTransaction);
         }
-    }
 
-    try {
         for (int tries = 0;; ++tries) {
             // Try kMaxNumStaleVersionRetries times. On the last try, exceptions are rethrown.
             bool canRetry = tries < kMaxNumStaleVersionRetries - 1;
@@ -423,6 +428,8 @@ void runCommand(OperationContext* opCtx,
     } catch (const DBException& e) {
         command->incrementCommandsFailed();
         LastError::get(opCtx->getClient()).setLastError(e.code(), e.reason());
+        auto errorLabels = getErrorLabels(osi, command->getName(), e.code());
+        errorBuilder->appendElements(errorLabels);
         throw;
     }
 }
@@ -534,6 +541,7 @@ DbResponse Strategy::queryOp(OperationContext* opCtx, const NamespaceString& nss
 
 DbResponse Strategy::clientCommand(OperationContext* opCtx, const Message& m) {
     auto reply = rpc::makeReplyBuilder(rpc::protocolForMessage(m));
+    BSONObjBuilder errorBuilder;
 
     bool propagateException = false;
 
@@ -556,7 +564,7 @@ DbResponse Strategy::clientCommand(OperationContext* opCtx, const Message& m) {
         std::string db = request.getDatabase().toString();
         try {
             LOG(3) << "Command begin db: " << db << " msg id: " << m.header().getId();
-            runCommand(opCtx, request, m.operation(), reply.get());
+            runCommand(opCtx, request, m.operation(), reply.get(), &errorBuilder);
             LOG(3) << "Command end db: " << db << " msg id: " << m.header().getId();
         } catch (const DBException& ex) {
             LOG(1) << "Exception thrown while processing command on " << db
@@ -574,6 +582,7 @@ DbResponse Strategy::clientCommand(OperationContext* opCtx, const Message& m) {
         auto bob = reply->getBodyBuilder();
         CommandHelpers::appendCommandStatusNoThrow(bob, ex.toStatus());
         appendRequiredFieldsToResponse(opCtx, &bob);
+        bob.appendElements(errorBuilder.obj());
     }
 
     if (OpMsg::isFlagSet(m, OpMsg::kMoreToCome)) {
@@ -718,6 +727,7 @@ void Strategy::killCursors(OperationContext* opCtx, DbMessage* dbm) {
 void Strategy::writeOp(OperationContext* opCtx, DbMessage* dbm) {
     const auto& msg = dbm->msg();
     rpc::OpMsgReplyBuilder reply;
+    BSONObjBuilder errorBuilder;
     runCommand(opCtx,
                [&]() {
                    switch (msg.operation()) {
@@ -735,7 +745,8 @@ void Strategy::writeOp(OperationContext* opCtx, DbMessage* dbm) {
                    }
                }(),
                msg.operation(),
-               &reply);  // built object is ignored
+               &reply,
+               &errorBuilder);  // built objects are ignored
 }
 
 void Strategy::explainFind(OperationContext* opCtx,