From 9b8d4c2f029e2384fa769014c00091e4de362ea4 Mon Sep 17 00:00:00 2001
From: "A. Jesse Jiryu Davis" <jesse@mongodb.com>
Date: Fri, 2 Oct 2020 11:11:45 -0400
Subject: SERVER-50375 Ensure mongos forwards API parameters to shards

---
 src/mongo/db/session_killer.cpp | 44 ++++++++++++++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 12 deletions(-)

(limited to 'src/mongo/db/session_killer.cpp')

diff --git a/src/mongo/db/session_killer.cpp b/src/mongo/db/session_killer.cpp
index 20ba36ac3dd..8ac4764398a 100644
--- a/src/mongo/db/session_killer.cpp
+++ b/src/mongo/db/session_killer.cpp
@@ -85,18 +85,19 @@ SessionKiller::ReapResult::ReapResult() : result(std::make_shared<boost::optiona
 
 SessionKiller::Matcher::Matcher(KillAllSessionsByPatternSet&& patterns)
     : _patterns(std::move(patterns)) {
-    for (const auto& pattern : _patterns) {
+    for (const auto& item : _patterns) {
+        auto& pattern = item.pattern;
         if (pattern.getUid()) {
             _uids.emplace(pattern.getUid().get(), &pattern);
         } else if (pattern.getLsid()) {
             _lsids.emplace(pattern.getLsid().get(), &pattern);
         } else {
             // If we're killing everything, it's the only pattern we care about.
-            decltype(_patterns) onlyKillAll{{pattern}};
+            decltype(_patterns) onlyKillAll{{item}};
             using std::swap;
             swap(_patterns, onlyKillAll);
 
-            _killAll = &(*_patterns.begin());
+            _killAll = &(_patterns.begin()->pattern);
             break;
         }
     }
@@ -177,19 +178,38 @@ void SessionKiller::_periodicKill(OperationContext* opCtx, stdx::unique_lock<Lat
     // Drop the lock and run the killer.
     lk.unlock();
 
-    Matcher matcher(std::move(nextToReap));
-    boost::optional<Result> results;
-    try {
-        results.emplace(_killFunc(opCtx, matcher, &_urbg));
-    } catch (...) {
-        results.emplace(exceptionToStatus());
+    // Group patterns with equal API parameters into sets.
+    stdx::unordered_map<APIParameters, KillAllSessionsByPatternSet, APIParameters::Hash> sets;
+    for (auto& item : nextToReap) {
+        sets[item.apiParameters].insert(std::move(item));
     }
-    lk.lock();
 
-    invariant(results);
+    // Use the API parameters included in each KillAllSessionsByPattern struct.
+    IgnoreAPIParametersBlock ignoreApiParametersBlock(opCtx);
+    Result finalResults{std::vector<HostAndPort>{}};
+    for (auto& [apiParameters, items] : sets) {
+        APIParameters::get(opCtx) = apiParameters;
+        Matcher matcher(std::move(items));
+        boost::optional<Result> results;
+        try {
+            results.emplace(_killFunc(opCtx, matcher, &_urbg));
+        } catch (...) {
+            results.emplace(exceptionToStatus());
+        }
+        invariant(results);
+        if (!results->isOK()) {
+            finalResults = *results;
+            break;
+        }
+
+        finalResults.getValue().insert(finalResults.getValue().end(),
+                                       std::make_move_iterator(results->getValue().begin()),
+                                       std::make_move_iterator(results->getValue().end()));
+    }
 
     // Expose the results and notify callers
-    *(reapResults.result) = std::move(results);
+    lk.lock();
+    *(reapResults.result) = std::move(finalResults);
     _callerCV.notify_all();
 };
 
-- 
cgit v1.2.1