diff options
Diffstat (limited to 'src/mongo/executor')
30 files changed, 3574 insertions, 563 deletions
diff --git a/src/mongo/executor/SConscript b/src/mongo/executor/SConscript index 627675845a2..355a0d6cc06 100644 --- a/src/mongo/executor/SConscript +++ b/src/mongo/executor/SConscript @@ -119,6 +119,16 @@ env.Library( ) env.Library( + target='connection_pool_controllers', + source=[ + 'connection_pool_controllers.cpp', + ], + LIBDEPS=[ + 'connection_pool_executor', + ], +) + +env.Library( target='network_test_env', source=[ 'network_test_env.cpp', @@ -231,6 +241,20 @@ env.Library( ) env.Library( + target='pinned_connection_task_executor', + source=[ + 'pinned_connection_task_executor.cpp', + ], + LIBDEPS=[ + 'scoped_task_executor', + 'task_executor_interface', + ], + LIBDEPS_PRIVATE=[ + '$BUILD_DIR/mongo/client/async_client', + ], +) + +env.Library( target='network_interface_thread_pool', source=[ 'network_interface_thread_pool.cpp', @@ -268,11 +292,15 @@ env.Library( target='task_executor_cursor', source=[ 'task_executor_cursor.cpp', + 'task_executor_cursor_parameters.idl', ], LIBDEPS=[ '$BUILD_DIR/mongo/db/query/command_request_response', 'task_executor_interface', ], + LIBDEPS_PRIVATE=[ + 'pinned_connection_task_executor_factory', + ], ) env.Library( @@ -288,6 +316,21 @@ env.Library( ], ) +env.Library( + target='pinned_connection_task_executor_factory', + source=[ + 'pinned_connection_task_executor_factory.cpp', + ], + LIBDEPS=[ + 'task_executor_interface', + ], + LIBDEPS_PRIVATE=[ + 'network_interface', + 'pinned_connection_task_executor', + 'thread_pool_task_executor', + ], +) + env.CppUnitTest( target='executor_test', source=[ @@ -297,14 +340,17 @@ env.CppUnitTest( 'mock_network_fixture_test.cpp', 'network_interface_mock_test.cpp', 'network_interface_mock_test_fixture.cpp', + 'pinned_connection_task_executor_test.cpp', 'scoped_task_executor_test.cpp', 'task_executor_cursor_test.cpp', 'thread_pool_task_executor_test.cpp', ], LIBDEPS=[ + '$BUILD_DIR/mongo/transport/message_compressor', 'connection_pool_executor', 'egress_tag_closer_manager', 'network_interface_mock', + 'pinned_connection_task_executor', 'scoped_task_executor', 'task_executor_cursor', 'thread_pool_task_executor', @@ -320,10 +366,12 @@ env.CppIntegrationTest( 'thread_pool_task_executor_integration_test.cpp', ], LIBDEPS=[ + '$BUILD_DIR/mongo/client/async_client', '$BUILD_DIR/mongo/client/clientdriver_network', '$BUILD_DIR/mongo/db/wire_version', '$BUILD_DIR/mongo/executor/network_interface_factory', '$BUILD_DIR/mongo/executor/network_interface_thread_pool', + '$BUILD_DIR/mongo/executor/pinned_connection_task_executor', '$BUILD_DIR/mongo/executor/thread_pool_task_executor', '$BUILD_DIR/mongo/transport/transport_layer_egress_init', '$BUILD_DIR/mongo/util/concurrency/thread_pool', diff --git a/src/mongo/executor/connection_pool.cpp b/src/mongo/executor/connection_pool.cpp index 5900f65a650..91fb862bb90 100644 --- a/src/mongo/executor/connection_pool.cpp +++ b/src/mongo/executor/connection_pool.cpp @@ -137,8 +137,8 @@ std::string ConnectionPool::ConnectionControls::toString() const { } std::string ConnectionPool::HostState::toString() const { - return "{{ requests: {}, ready: {}, pending: {}, active: {}, isExpired: {} }}"_format( - requests, ready, pending, active, health.isExpired); + return "{{ requests: {}, ready: {}, pending: {}, active: {}, leased: {}, isExpired: {} }}"_format( + requests, ready, pending, active, leased, health.isExpired); } /** @@ -162,7 +162,7 @@ public: const auto minConns = getPool()->_options.minConnections; const auto maxConns = getPool()->_options.maxConnections; - data.target = stats.requests + stats.active; + data.target = stats.requests + stats.active + stats.leased; if (data.target < minConns) { data.target = minConns; } else if (data.target > maxConns) { @@ -275,7 +275,7 @@ public: * Gets a connection from the specific pool. Sinks a unique_lock from the * parent to preserve the lock on _mutex */ - Future<ConnectionHandle> getConnection(Milliseconds timeout); + Future<ConnectionHandle> getConnection(Milliseconds timeout, bool lease); /** * Triggers the shutdown procedure. This function sets isShutdown to true @@ -298,6 +298,11 @@ public: size_t inUseConnections() const; /** + * Returns the number of leased connections from the pool. + */ + size_t leasedConnections() const; + + /** * Returns the number of available connections in the pool. */ size_t availableConnections() const; @@ -320,7 +325,7 @@ public: /** * Returns the total number of connections currently open that belong to * this pool. This is the sum of refreshingConnections, availableConnections, - * and inUseConnections. + * inUseConnections, and leasedConnections. */ size_t openConnections() const; @@ -361,14 +366,20 @@ private: using OwnedConnection = std::shared_ptr<ConnectionInterface>; using OwnershipPool = stdx::unordered_map<ConnectionInterface*, OwnedConnection>; using LRUOwnershipPool = LRUCache<OwnershipPool::key_type, OwnershipPool::mapped_type>; - using Request = std::pair<Date_t, Promise<ConnectionHandle>>; + struct Request { + Date_t expiration; + Promise<ConnectionHandle> promise; + // Whether or not the requested connection should be "leased". + bool lease; + }; + struct RequestComparator { bool operator()(const Request& a, const Request& b) { - return a.first > b.first; + return a.expiration > b.expiration; } }; - ConnectionHandle makeHandle(ConnectionInterface* connection); + ConnectionHandle makeHandle(ConnectionInterface* connection, bool isLeased); /** * Establishes connections until the ControllerInterface's target is met. @@ -381,11 +392,11 @@ private: void fulfillRequests(); - void returnConnection(ConnectionInterface* connPtr); + void returnConnection(ConnectionInterface* connPtr, bool isLeased); // This internal helper is used both by get and by _fulfillRequests and differs in that it // skips some bookkeeping that the other callers do on their own - ConnectionHandle tryGetConnection(); + ConnectionHandle tryGetConnection(bool lease); template <typename OwnershipPoolType> typename OwnershipPoolType::mapped_type takeFromPool( @@ -414,6 +425,7 @@ private: OwnershipPool _processingPool; OwnershipPool _droppedProcessingPool; OwnershipPool _checkedOutPool; + OwnershipPool _leasedPool; std::vector<Request> _requests; Date_t _lastActiveTime; @@ -548,21 +560,37 @@ void ConnectionPool::mutateTags( pool->mutateTags(mutateFunc); } +void ConnectionPool::retrieve_forTest(RetrieveConnection retrieve, GetConnectionCallback cb) { + // We kick ourselves onto the executor queue to prevent us from deadlocking with our own thread + auto getConnectionFunc = + [this, retrieve = std::move(retrieve), cb = std::move(cb)](Status&&) mutable { + retrieve().thenRunOn(_factory->getExecutor()).getAsync(std::move(cb)); + }; + _factory->getExecutor()->schedule(std::move(getConnectionFunc)); +} + void ConnectionPool::get_forTest(const HostAndPort& hostAndPort, Milliseconds timeout, GetConnectionCallback cb) { - // We kick ourselves onto the executor queue to prevent us from deadlocking with our own thread - auto getConnectionFunc = [this, hostAndPort, timeout, cb = std::move(cb)](Status&&) mutable { - get(hostAndPort, transport::kGlobalSSLMode, timeout) - .thenRunOn(_factory->getExecutor()) - .getAsync(std::move(cb)); + auto getConnectionFunc = [this, hostAndPort, timeout]() mutable { + return get(hostAndPort, transport::kGlobalSSLMode, timeout); }; - _factory->getExecutor()->schedule(std::move(getConnectionFunc)); + retrieve_forTest(getConnectionFunc, std::move(cb)); +} + +void ConnectionPool::lease_forTest(const HostAndPort& hostAndPort, + Milliseconds timeout, + GetConnectionCallback cb) { + auto getConnectionFunc = [this, hostAndPort, timeout]() mutable { + return lease(hostAndPort, transport::kGlobalSSLMode, timeout); + }; + retrieve_forTest(getConnectionFunc, std::move(cb)); } -SemiFuture<ConnectionPool::ConnectionHandle> ConnectionPool::get(const HostAndPort& hostAndPort, - transport::ConnectSSLMode sslMode, - Milliseconds timeout) { +SemiFuture<ConnectionPool::ConnectionHandle> ConnectionPool::_get(const HostAndPort& hostAndPort, + transport::ConnectSSLMode sslMode, + Milliseconds timeout, + bool lease) { stdx::lock_guard lk(_mutex); auto& pool = _pools[hostAndPort]; @@ -574,7 +602,7 @@ SemiFuture<ConnectionPool::ConnectionHandle> ConnectionPool::get(const HostAndPo invariant(pool); - auto connFuture = pool->getConnection(timeout); + auto connFuture = pool->getConnection(timeout, lease); pool->updateState(); return std::move(connFuture).semi(); @@ -590,6 +618,7 @@ void ConnectionPool::appendConnectionStats(ConnectionPoolStats* stats) const { auto& pool = kv.second; ConnectionStatsPer hostStats{pool->inUseConnections(), pool->availableConnections(), + pool->leasedConnections(), pool->createdConnections(), pool->refreshingConnections(), pool->refreshedConnections()}; @@ -625,6 +654,7 @@ ConnectionPool::SpecificPool::~SpecificPool() { if (shouldInvariantOnPoolCorrectness()) { invariant(_requests.empty()); invariant(_checkedOutPool.empty()); + invariant(_leasedPool.empty()); } } @@ -636,6 +666,10 @@ size_t ConnectionPool::SpecificPool::availableConnections() const { return _readyPool.size(); } +size_t ConnectionPool::SpecificPool::leasedConnections() const { + return _leasedPool.size(); +} + size_t ConnectionPool::SpecificPool::refreshingConnections() const { return _processingPool.size(); } @@ -649,7 +683,7 @@ size_t ConnectionPool::SpecificPool::createdConnections() const { } size_t ConnectionPool::SpecificPool::openConnections() const { - return _checkedOutPool.size() + _readyPool.size() + _processingPool.size(); + return _checkedOutPool.size() + _readyPool.size() + _processingPool.size() + _leasedPool.size(); } size_t ConnectionPool::SpecificPool::requestsPending() const { @@ -657,7 +691,7 @@ size_t ConnectionPool::SpecificPool::requestsPending() const { } Future<ConnectionPool::ConnectionHandle> ConnectionPool::SpecificPool::getConnection( - Milliseconds timeout) { + Milliseconds timeout, bool lease) { // Reset our activity timestamp auto now = _parent->_factory->now(); @@ -665,7 +699,7 @@ Future<ConnectionPool::ConnectionHandle> ConnectionPool::SpecificPool::getConnec // If we do not have requests, then we can fulfill immediately if (_requests.size() == 0) { - auto conn = tryGetConnection(); + auto conn = tryGetConnection(lease); if (conn) { LOGV2_DEBUG(22559, @@ -691,23 +725,24 @@ Future<ConnectionPool::ConnectionHandle> ConnectionPool::SpecificPool::getConnec const auto expiration = now + timeout; auto pf = makePromiseFuture<ConnectionHandle>(); - _requests.push_back(make_pair(expiration, std::move(pf.promise))); + _requests.push_back({expiration, std::move(pf.promise), lease}); std::push_heap(begin(_requests), end(_requests), RequestComparator{}); return std::move(pf.future); } -auto ConnectionPool::SpecificPool::makeHandle(ConnectionInterface* connection) -> ConnectionHandle { - auto deleter = [this, anchor = shared_from_this()](ConnectionInterface* connection) { +auto ConnectionPool::SpecificPool::makeHandle(ConnectionInterface* connection, bool isLeased) + -> ConnectionHandle { + auto deleter = [this, anchor = shared_from_this(), isLeased](ConnectionInterface* connection) { stdx::lock_guard lk(_parent->_mutex); - returnConnection(connection); + returnConnection(connection, isLeased); _lastActiveTime = _parent->_factory->now(); updateState(); }; return ConnectionHandle(connection, std::move(deleter)); } -ConnectionPool::ConnectionHandle ConnectionPool::SpecificPool::tryGetConnection() { +ConnectionPool::ConnectionHandle ConnectionPool::SpecificPool::tryGetConnection(bool lease) { while (_readyPool.size()) { // _readyPool is an LRUCache, so its begin() object is the MRU item. auto iter = _readyPool.begin(); @@ -729,12 +764,15 @@ ConnectionPool::ConnectionHandle ConnectionPool::SpecificPool::tryGetConnection( auto connPtr = conn.get(); - // check out the connection - _checkedOutPool[connPtr] = std::move(conn); + if (lease) { + _leasedPool[connPtr] = std::move(conn); + } else { + _checkedOutPool[connPtr] = std::move(conn); + } // pass it to the user connPtr->resetToUnknown(); - auto handle = makeHandle(connPtr); + auto handle = makeHandle(connPtr, lease); return handle; } @@ -770,15 +808,14 @@ void ConnectionPool::SpecificPool::finishRefresh(ConnectionInterface* connPtr, S return; } - // Pass a failure on through - if (!status.isOK()) { - LOGV2_DEBUG(22563, + // If the error can be contained to one connection, drop the one connection. + if (status.code() == ErrorCodes::ConnectionError) { + LOGV2_DEBUG(6832901, kDiagnosticLogLevel, - "Connection failed to {hostAndPort} due to {error}", - "Connection failed", + "Dropping single connection", "hostAndPort"_attr = _hostAndPort, - "error"_attr = redact(status)); - processFailure(status); + "error"_attr = redact(status), + "numOpenConns"_attr = openConnections()); return; } @@ -792,6 +829,18 @@ void ConnectionPool::SpecificPool::finishRefresh(ConnectionInterface* connPtr, S return; } + // Pass a failure on through + if (!status.isOK()) { + LOGV2_DEBUG(22563, + kDiagnosticLogLevel, + "Connection failed to {hostAndPort} due to {error}", + "Connection failed", + "hostAndPort"_attr = _hostAndPort, + "error"_attr = redact(status)); + processFailure(status); + return; + } + LOGV2_DEBUG(22565, kDiagnosticLogLevel, "Finishing connection refresh for {hostAndPort}", @@ -804,10 +853,10 @@ void ConnectionPool::SpecificPool::finishRefresh(ConnectionInterface* connPtr, S fulfillRequests(); } -void ConnectionPool::SpecificPool::returnConnection(ConnectionInterface* connPtr) { +void ConnectionPool::SpecificPool::returnConnection(ConnectionInterface* connPtr, bool isLeased) { auto needsRefreshTP = connPtr->getLastUsed() + _parent->_controller->toRefreshTimeout(); - auto conn = takeFromPool(_checkedOutPool, connPtr); + auto conn = takeFromPool(isLeased ? _leasedPool : _checkedOutPool, connPtr); invariant(conn); if (_health.isShutdown) { @@ -821,7 +870,29 @@ void ConnectionPool::SpecificPool::returnConnection(ConnectionInterface* connPtr } if (auto status = conn->getStatus(); !status.isOK()) { - // TODO: alert via some callback if the host is bad + // Our error handling here is determined by the MongoDB SDAM specification for handling + // application errors on established connections. In particular, if a network error occurs, + // we must close all idle sockets in the connection pool for the server: "if one socket is + // bad, it is likely that all are." However, if the error is just a network _timeout_ error, + // we don't drop the connections because the timeout may indicate a slow operation rather + // than an unavailable server. Additionally, if we can isolate the error to a single + // socket/connection based on it's type, we won't drop other connections/sockets. + // + // See the spec for additional details: + // https://github.com/mongodb/specifications/blob/master/source/server-discovery-and-monitoring/server-discovery-and-monitoring.rst#application-errors + bool isSingleConnectionError = status.code() == ErrorCodes::ConnectionError; + if (ErrorCodes::isNetworkError(status) && !isSingleConnectionError && + !ErrorCodes::isNetworkTimeoutError(status)) { + LOGV2_DEBUG(7719500, + kDiagnosticLogLevel, + "Connection failed to {hostAndPort} due to {error}", + "Connection failed", + "hostAndPort"_attr = _hostAndPort, + "error"_attr = redact(status)); + processFailure(status); + return; + } + // Otherwise, drop the one connection. LOGV2(22566, "Ending connection to host {hostAndPort} due to bad connection status: {error}; " "{numOpenConns} connections to that host remain open", @@ -842,8 +913,7 @@ void ConnectionPool::SpecificPool::returnConnection(ConnectionInterface* connPtr if (shouldRefreshConnection) { auto controls = _parent->_controller->getControls(_id); - if (_readyPool.size() + _processingPool.size() + _checkedOutPool.size() >= - controls.targetConnections) { + if (openConnections() >= controls.targetConnections) { // If we already have minConnections, just let the connection lapse LOGV2(22567, "Ending idle connection to host {hostAndPort} because the pool meets " @@ -910,7 +980,7 @@ void ConnectionPool::SpecificPool::addToReady(OwnedConnection conn) { connPtr->indicateSuccess(); - returnConnection(connPtr); + returnConnection(connPtr, false); }); connPtr->setTimeout(_parent->_controller->toRefreshTimeout(), std::move(returnConnectionFunc)); } @@ -977,7 +1047,7 @@ void ConnectionPool::SpecificPool::processFailure(const Status& status) { } for (auto& request : _requests) { - request.second.setError(status); + request.promise.setError(status); } LOGV2_DEBUG(22573, @@ -999,14 +1069,14 @@ void ConnectionPool::SpecificPool::fulfillRequests() { // deadlock). // // None of the heap manipulation code throws, but it's something to keep in mind. - auto conn = tryGetConnection(); + auto conn = tryGetConnection(_requests.front().lease); if (!conn) { break; } // Grab the request and callback - auto promise = std::move(_requests.front().second); + auto promise = std::move(_requests.front().promise); std::pop_heap(begin(_requests), end(_requests), RequestComparator{}); _requests.pop_back(); @@ -1114,7 +1184,8 @@ void ConnectionPool::SpecificPool::updateHealth() { const auto now = _parent->_factory->now(); // We're expired if we have no sign of connection use and are past our expiry - _health.isExpired = _requests.empty() && _checkedOutPool.empty() && (_hostExpiration <= now); + _health.isExpired = _requests.empty() && _checkedOutPool.empty() && _leasedPool.empty() && + (_hostExpiration <= now); // We're failed until we get new requests or our timer triggers if (_health.isFailed) { @@ -1132,7 +1203,7 @@ void ConnectionPool::SpecificPool::updateEventTimer() { } // If our expiration comes before our next event, then it is the next event - if (_requests.empty() && _checkedOutPool.empty()) { + if (_requests.empty() && _checkedOutPool.empty() && _leasedPool.empty()) { _hostExpiration = _lastActiveTime + _parent->_controller->hostTimeout(); if ((_hostExpiration > now) && (_hostExpiration < nextEventTime)) { nextEventTime = _hostExpiration; @@ -1140,8 +1211,8 @@ void ConnectionPool::SpecificPool::updateEventTimer() { } // If a request would timeout before the next event, then it is the next event - if (_requests.size() && (_requests.front().first < nextEventTime)) { - nextEventTime = _requests.front().first; + if (_requests.size() && (_requests.front().expiration < nextEventTime)) { + nextEventTime = _requests.front().expiration; } // Clamp next event time to be either now or in the future. Next event time @@ -1169,12 +1240,12 @@ void ConnectionPool::SpecificPool::updateEventTimer() { _health.isFailed = false; - while (_requests.size() && (_requests.front().first <= now)) { + while (_requests.size() && (_requests.front().expiration <= now)) { std::pop_heap(begin(_requests), end(_requests), RequestComparator{}); auto& request = _requests.back(); - request.second.setError(Status(ErrorCodes::NetworkInterfaceExceededTimeLimit, - "Couldn't get a connection within the time limit")); + request.promise.setError(Status(ErrorCodes::NetworkInterfaceExceededTimeLimit, + "Couldn't get a connection within the time limit")); _requests.pop_back(); // Since we've failed a request, we've interacted with external users @@ -1198,6 +1269,7 @@ void ConnectionPool::SpecificPool::updateController() { refreshingConnections(), availableConnections(), inUseConnections(), + leasedConnections(), }; LOGV2_DEBUG(22578, kDiagnosticLogLevel, @@ -1236,6 +1308,7 @@ void ConnectionPool::SpecificPool::updateController() { if (shouldInvariantOnPoolCorrectness()) { invariant(pool->_checkedOutPool.empty()); invariant(pool->_requests.empty()); + invariant(pool->_leasedPool.empty()); } pool->triggerShutdown(Status(ErrorCodes::ConnectionPoolExpired, diff --git a/src/mongo/executor/connection_pool.h b/src/mongo/executor/connection_pool.h index 0e5bf90dc9a..b815003f4d8 100644 --- a/src/mongo/executor/connection_pool.h +++ b/src/mongo/executor/connection_pool.h @@ -79,6 +79,7 @@ public: using ConnectionHandleDeleter = std::function<void(ConnectionInterface* connection)>; using ConnectionHandle = std::unique_ptr<ConnectionInterface, ConnectionHandleDeleter>; + using RetrieveConnection = unique_function<SemiFuture<ConnectionHandle>()>; using GetConnectionCallback = unique_function<void(StatusWith<ConnectionHandle>)>; using PoolId = uint64_t; @@ -210,6 +211,7 @@ public: size_t pending = 0; size_t ready = 0; size_t active = 0; + size_t leased = 0; std::string toString() const; }; @@ -252,13 +254,35 @@ public: const std::function<transport::Session::TagMask(transport::Session::TagMask)>& mutateFunc) override; - SemiFuture<ConnectionHandle> get(const HostAndPort& hostAndPort, - transport::ConnectSSLMode sslMode, - Milliseconds timeout); + inline SemiFuture<ConnectionHandle> get(const HostAndPort& hostAndPort, + transport::ConnectSSLMode sslMode, + Milliseconds timeout) { + return _get(hostAndPort, sslMode, timeout, false /*lease*/); + } + void get_forTest(const HostAndPort& hostAndPort, Milliseconds timeout, GetConnectionCallback cb); + /** + * "Lease" a connection from the pool. + * + * Connections retrieved via this method are not assumed to be in active use for the duration of + * their lease and are reported separately in metrics. Otherwise, this method behaves similarly + * to `ConnectionPool::get`. + */ + inline SemiFuture<ConnectionHandle> lease( + const HostAndPort& hostAndPort, + transport::ConnectSSLMode sslMode, + Milliseconds timeout, + ErrorCodes::Error timeoutCode = ErrorCodes::NetworkInterfaceExceededTimeLimit) { + return _get(hostAndPort, sslMode, timeout, true /*lease*/); + } + + void lease_forTest(const HostAndPort& hostAndPort, + Milliseconds timeout, + GetConnectionCallback cb); + void appendConnectionStats(ConnectionPoolStats* stats) const; size_t getNumConnectionsPerHost(const HostAndPort& hostAndPort) const; @@ -268,6 +292,13 @@ public: } private: + SemiFuture<ConnectionHandle> _get(const HostAndPort& hostAndPort, + transport::ConnectSSLMode sslMode, + Milliseconds timeout, + bool leased); + + void retrieve_forTest(RetrieveConnection retrieve, GetConnectionCallback cb); + std::string _name; const std::shared_ptr<DependentTypeFactoryInterface> _factory; @@ -496,6 +527,10 @@ public: return _pool; } + Options getPoolOptions() const { + return _pool->_options; + } + virtual void updateConnectionPoolStats([[maybe_unused]] ConnectionPoolStats* cps) const = 0; protected: diff --git a/src/mongo/executor/connection_pool_controllers.cpp b/src/mongo/executor/connection_pool_controllers.cpp new file mode 100644 index 00000000000..10864bcbd03 --- /dev/null +++ b/src/mongo/executor/connection_pool_controllers.cpp @@ -0,0 +1,78 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 <algorithm> + +#include "mongo/executor/connection_pool_controllers.h" + +namespace mongo::executor { +namespace { +template <typename Map, typename Key> +auto& getOrInvariant(Map&& map, const Key& key) { + auto it = map.find(key); + invariant(it != map.end(), "Unable to find key in map"); + + return it->second; +} +} // namespace + +void DynamicLimitController::init(executor::ConnectionPool* parent) { + ControllerInterface::init(parent); +} + +void DynamicLimitController::addHost(PoolId id, const HostAndPort& host) { + stdx::lock_guard lk(_mutex); + auto ret = _poolData.insert({id, {host}}); + using namespace fmt::literals; + invariant( + ret.second, + "ConnectionPool controller {} received a request to track host {} that was already being tracked."_format( + _name, host)); +} + +DynamicLimitController::HostGroupState DynamicLimitController::updateHost(PoolId id, + const HostState& stats) { + stdx::lock_guard lk(_mutex); + auto& data = getOrInvariant(_poolData, id); + data.target = + std::clamp(stats.requests + stats.active + stats.leased, _minLoader(), _maxLoader()); + return {{data.host}, stats.health.isExpired}; +} + +void DynamicLimitController::removeHost(PoolId id) { + stdx::lock_guard lk(_mutex); + invariant(_poolData.erase(id)); +} + +ConnectionPool::ConnectionControls DynamicLimitController::getControls(PoolId id) { + stdx::lock_guard lk(_mutex); + return {getPoolOptions().maxConnecting, getOrInvariant(_poolData, id).target}; +} + +} // namespace mongo::executor diff --git a/src/mongo/executor/connection_pool_controllers.h b/src/mongo/executor/connection_pool_controllers.h new file mode 100644 index 00000000000..9e7929ec884 --- /dev/null +++ b/src/mongo/executor/connection_pool_controllers.h @@ -0,0 +1,107 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 <functional> +#include <string> + +#include "mongo/base/string_data.h" +#include "mongo/executor/connection_pool.h" +#include "mongo/executor/connection_pool_stats.h" +#include "mongo/platform/mutex.h" +#include "mongo/stdx/unordered_map.h" +#include "mongo/util/duration.h" +#include "mongo/util/net/hostandport.h" + +namespace mongo::executor { +/** + * This file is intended for simple implementations of ConnectionPool::ControllerInterface that + * might be shared between different libraries. Currently, it contains only one such implementation, + * the DyamicLimitController below. + */ + +/** + * A simple controller that allows for the maximum and minimum pool size to have dynamic values. + * At construction, provide callables that return the current maximum and minimum sizes to be used + * by the pool. + * + * Currently, the callables that provide the max and min are stateless and don't inspect any data + * about the pool. The other pool parameters (maxConnecting, host/pending/refresh timeouts) are + * simply taken from the Optoins the relevant ConnectionPool was started with. However, this + * type is intended to be easily extensible to add these features in the future if needed. + */ +class DynamicLimitController final : public ConnectionPool::ControllerInterface { +public: + DynamicLimitController(std::function<size_t()> minLoader, + std::function<size_t()> maxLoader, + StringData name) + : _minLoader(std::move(minLoader)), + _maxLoader(std::move(maxLoader)), + _name(std::move(name)) {} + + void init(ConnectionPool* parent) override; + + void addHost(PoolId id, const HostAndPort& host) override; + HostGroupState updateHost(PoolId id, const HostState& stats) override; + void removeHost(PoolId id) override; + + ConnectionControls getControls(PoolId id) override; + + Milliseconds hostTimeout() const override { + return getPoolOptions().hostTimeout; + } + + Milliseconds pendingTimeout() const override { + return getPoolOptions().refreshTimeout; + } + + Milliseconds toRefreshTimeout() const override { + return getPoolOptions().refreshRequirement; + } + + StringData name() const override { + return _name; + } + + void updateConnectionPoolStats(ConnectionPoolStats* cps) const override {} + +private: + struct PoolData { + HostAndPort host; + size_t target = 0; + }; + + std::function<size_t()> _minLoader; + std::function<size_t()> _maxLoader; + std::string _name; + Mutex _mutex = MONGO_MAKE_LATCH("DynamicLimitController::_mutex"); + stdx::unordered_map<PoolId, PoolData> _poolData; +}; +} // namespace mongo::executor diff --git a/src/mongo/executor/connection_pool_stats.cpp b/src/mongo/executor/connection_pool_stats.cpp index a3d1e1fb0b1..85d6743b127 100644 --- a/src/mongo/executor/connection_pool_stats.cpp +++ b/src/mongo/executor/connection_pool_stats.cpp @@ -36,10 +36,15 @@ namespace mongo { namespace executor { -ConnectionStatsPer::ConnectionStatsPer( - size_t nInUse, size_t nAvailable, size_t nCreated, size_t nRefreshing, size_t nRefreshed) +ConnectionStatsPer::ConnectionStatsPer(size_t nInUse, + size_t nAvailable, + size_t nLeased, + size_t nCreated, + size_t nRefreshing, + size_t nRefreshed) : inUse(nInUse), available(nAvailable), + leased(nLeased), created(nCreated), refreshing(nRefreshing), refreshed(nRefreshed) {} @@ -49,6 +54,7 @@ ConnectionStatsPer::ConnectionStatsPer() = default; ConnectionStatsPer& ConnectionStatsPer::operator+=(const ConnectionStatsPer& other) { inUse += other.inUse; available += other.available; + leased += other.leased; created += other.created; refreshing += other.refreshing; refreshed += other.refreshed; @@ -75,6 +81,7 @@ void ConnectionPoolStats::updateStatsForHost(std::string pool, // Update total connection stats. totalInUse += newStats.inUse; totalAvailable += newStats.available; + totalLeased += newStats.leased; totalCreated += newStats.created; totalRefreshing += newStats.refreshing; totalRefreshed += newStats.refreshed; @@ -83,6 +90,7 @@ void ConnectionPoolStats::updateStatsForHost(std::string pool, void ConnectionPoolStats::appendToBSON(mongo::BSONObjBuilder& result, bool forFTDC) { result.appendNumber("totalInUse", static_cast<long long>(totalInUse)); result.appendNumber("totalAvailable", static_cast<long long>(totalAvailable)); + result.appendNumber("totalLeased", static_cast<long long>(totalLeased)); result.appendNumber("totalCreated", static_cast<long long>(totalCreated)); result.appendNumber("totalRefreshing", static_cast<long long>(totalRefreshing)); result.appendNumber("totalRefreshed", static_cast<long long>(totalRefreshed)); @@ -115,6 +123,7 @@ void ConnectionPoolStats::appendToBSON(mongo::BSONObjBuilder& result, bool forFT auto& poolStats = pool.second; poolInfo.appendNumber("poolInUse", static_cast<long long>(poolStats.inUse)); poolInfo.appendNumber("poolAvailable", static_cast<long long>(poolStats.available)); + poolInfo.appendNumber("poolLeased", static_cast<long long>(poolStats.leased)); poolInfo.appendNumber("poolCreated", static_cast<long long>(poolStats.created)); poolInfo.appendNumber("poolRefreshing", static_cast<long long>(poolStats.refreshing)); poolInfo.appendNumber("poolRefreshed", static_cast<long long>(poolStats.refreshed)); @@ -124,6 +133,7 @@ void ConnectionPoolStats::appendToBSON(mongo::BSONObjBuilder& result, bool forFT auto& hostStats = host.second; hostInfo.appendNumber("inUse", static_cast<long long>(hostStats.inUse)); hostInfo.appendNumber("available", static_cast<long long>(hostStats.available)); + hostInfo.appendNumber("leased", static_cast<long long>(hostStats.leased)); hostInfo.appendNumber("created", static_cast<long long>(hostStats.created)); hostInfo.appendNumber("refreshing", static_cast<long long>(hostStats.refreshing)); hostInfo.appendNumber("refreshed", static_cast<long long>(hostStats.refreshed)); @@ -139,6 +149,7 @@ void ConnectionPoolStats::appendToBSON(mongo::BSONObjBuilder& result, bool forFT auto hostStats = host.second; hostInfo.appendNumber("inUse", static_cast<long long>(hostStats.inUse)); hostInfo.appendNumber("available", static_cast<long long>(hostStats.available)); + hostInfo.appendNumber("leased", static_cast<long long>(hostStats.leased)); hostInfo.appendNumber("created", static_cast<long long>(hostStats.created)); hostInfo.appendNumber("refreshing", static_cast<long long>(hostStats.refreshing)); hostInfo.appendNumber("refreshed", static_cast<long long>(hostStats.refreshed)); diff --git a/src/mongo/executor/connection_pool_stats.h b/src/mongo/executor/connection_pool_stats.h index a96739b42be..3f4183adbf2 100644 --- a/src/mongo/executor/connection_pool_stats.h +++ b/src/mongo/executor/connection_pool_stats.h @@ -41,8 +41,12 @@ namespace executor { * a parent ConnectionPoolStats object and should not need to be created directly. */ struct ConnectionStatsPer { - ConnectionStatsPer( - size_t nInUse, size_t nAvailable, size_t nCreated, size_t nRefreshing, size_t nRefreshed); + ConnectionStatsPer(size_t nInUse, + size_t nAvailable, + size_t nLeased, + size_t nCreated, + size_t nRefreshing, + size_t nRefreshed); ConnectionStatsPer(); @@ -50,6 +54,7 @@ struct ConnectionStatsPer { size_t inUse = 0u; size_t available = 0u; + size_t leased = 0u; size_t created = 0u; size_t refreshing = 0u; size_t refreshed = 0u; @@ -68,6 +73,7 @@ struct ConnectionPoolStats { size_t totalInUse = 0u; size_t totalAvailable = 0u; + size_t totalLeased = 0u; size_t totalCreated = 0u; size_t totalRefreshing = 0u; size_t totalRefreshed = 0u; diff --git a/src/mongo/executor/connection_pool_test.cpp b/src/mongo/executor/connection_pool_test.cpp index 79eaaaf0218..cfc514e562e 100644 --- a/src/mongo/executor/connection_pool_test.cpp +++ b/src/mongo/executor/connection_pool_test.cpp @@ -31,6 +31,8 @@ #include "mongo/executor/connection_pool_test_fixture.h" +#include "mongo/util/duration.h" +#include "mongo/util/net/hostandport.h" #include <algorithm> #include <memory> #include <random> @@ -41,6 +43,7 @@ #include <fmt/ostream.h> #include "mongo/executor/connection_pool.h" +#include "mongo/executor/connection_pool_stats.h" #include "mongo/stdx/future.h" #include "mongo/unittest/thread_assertion_monitor.h" #include "mongo/unittest/unittest.h" @@ -52,6 +55,8 @@ namespace connection_pool_test_details { class ConnectionPoolTest : public unittest::Test { public: + constexpr static Milliseconds kNoTimeout = Milliseconds{-1}; + protected: void setUp() override {} @@ -91,6 +96,12 @@ protected: ExecutorFuture(_executor).getAsync([conn = std::move(conn)](auto) {}); } + void doneWithError(ConnectionPool::ConnectionHandle& conn, Status error) { + dynamic_cast<ConnectionImpl*>(conn.get())->indicateFailure(error); + + ExecutorFuture(_executor).getAsync([conn = std::move(conn)](auto) {}); + } + using StatusWithConn = StatusWith<ConnectionPool::ConnectionHandle>; auto getId(const ConnectionPool::ConnectionHandle& conn) { @@ -148,6 +159,8 @@ TEST_F(ConnectionPoolTest, SameConn) { */ TEST_F(ConnectionPoolTest, ConnectionsAreAcquiredInMRUOrder) { auto pool = makePool(); + std::random_device rd; + std::mt19937 rng(rd()); // Obtain a set of connections constexpr size_t kSize = 100; @@ -170,17 +183,24 @@ TEST_F(ConnectionPoolTest, ConnectionsAreAcquiredInMRUOrder) { } }); + std::uniform_int_distribution<> dist{0, 1}; for (size_t i = 0; i != kSize; ++i) { ConnectionImpl::pushSetup(Status::OK()); - pool->get_forTest(HostAndPort(), - Milliseconds(5000), - [&](StatusWith<ConnectionPool::ConnectionHandle> swConn) { - monitors[i].exec([&]() { - ASSERT(swConn.isOK()); - connections.push_back(std::move(swConn.getValue())); - monitors[i].notifyDone(); - }); - }); + auto cb = [&](StatusWith<ConnectionPool::ConnectionHandle> swConn) { + monitors[i].exec([&]() { + ASSERT(swConn.isOK()); + connections.push_back(std::move(swConn.getValue())); + monitors[i].notifyDone(); + }); + }; + auto timeout = Milliseconds(5000); + + // Randomly lease or check out connection. + if (dist(rng)) { + pool->get_forTest(HostAndPort(), timeout, cb); + } else { + pool->lease_forTest(HostAndPort(), timeout, cb); + } } for (auto& monitor : monitors) { @@ -190,8 +210,6 @@ TEST_F(ConnectionPoolTest, ConnectionsAreAcquiredInMRUOrder) { ASSERT_EQ(connections.size(), kSize); // Shuffle them into a random order - std::random_device rd; - std::mt19937 rng(rd()); std::shuffle(connections.begin(), connections.end(), rng); // Return them to the pool in that random order, recording IDs in a stack @@ -211,18 +229,24 @@ TEST_F(ConnectionPoolTest, ConnectionsAreAcquiredInMRUOrder) { // as the IDs in the stack, since the pool returns them in MRU order. for (size_t i = 0; i != kSize; ++i) { ConnectionImpl::pushSetup(Status::OK()); - pool->get_forTest(HostAndPort(), - Milliseconds(5000), - [&](StatusWith<ConnectionPool::ConnectionHandle> swConn) { - monitors[i].exec([&]() { - ASSERT(swConn.isOK()); - const auto id = verifyAndGetId(swConn); - connections.push_back(std::move(swConn.getValue())); - ASSERT_EQ(id, ids.top()); - ids.pop(); - monitors[i].notifyDone(); - }); - }); + auto cb = [&](StatusWith<ConnectionPool::ConnectionHandle> swConn) { + monitors[i].exec([&]() { + ASSERT(swConn.isOK()); + const auto id = verifyAndGetId(swConn); + connections.push_back(std::move(swConn.getValue())); + ASSERT_EQ(id, ids.top()); + ids.pop(); + monitors[i].notifyDone(); + }); + }; + auto timeout = Milliseconds(5000); + + // Randomly lease or check out connection. + if (dist(rng)) { + pool->get_forTest(HostAndPort(), timeout, cb); + } else { + pool->lease_forTest(HostAndPort(), timeout, cb); + } } for (auto& monitor : monitors) { @@ -384,6 +408,155 @@ TEST_F(ConnectionPoolTest, FailedConnDifferentConn) { } /** + * Verify that a connection returned with an error indicating the remote + * is unavailable drops current generation connections to that remote. + */ +TEST_F(ConnectionPoolTest, FailedHostDropsConns) { + auto pool = makePool(); + + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), 0U); + + constexpr size_t kSize = 100; + std::vector<ConnectionPool::ConnectionHandle> connections; + std::vector<unittest::ThreadAssertionMonitor> monitors(kSize); + + // Ensure that no matter how we leave the test, we mark any + // checked out connections as OK before implicity returning them + // to the pool by destroying the 'connections' vector. Otherwise, + // this test would cause an invariant failure instead of a normal + // test failure if it fails, which would be confusing. + auto drainConnPool = [&] { + while (!connections.empty()) { + try { + ConnectionPool::ConnectionHandle conn = std::move(connections.back()); + connections.pop_back(); + doneWith(conn); + } catch (...) { + } + } + }; + const ScopeGuard guard(drainConnPool); + + auto now = Date_t::now(); + PoolImpl::setNow(now); + + // Check out kSize connections from the pool. + for (size_t i = 0; i != kSize; ++i) { + ConnectionImpl::pushSetup(Status::OK()); + auto cb = [&](StatusWith<ConnectionPool::ConnectionHandle> swConn) { + monitors[i].exec([&]() { + ASSERT(swConn.isOK()); + connections.push_back(std::move(swConn.getValue())); + monitors[i].notifyDone(); + }); + }; + auto timeout = Milliseconds(5000); + pool->get_forTest(HostAndPort(), timeout, cb); + } + + for (auto& monitor : monitors) { + monitor.wait(); + } + + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), kSize); + + // Return one connection with a network error. + ConnectionPool::ConnectionHandle conn = std::move(connections.back()); + connections.pop_back(); + doneWithError(conn, {ErrorCodes::HostUnreachable, "error"}); + + // We should still have all of the connections open, minus the one we just returned with an + // error. + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), kSize - 1); + + // Put the remaining connections back. + drainConnPool(); + + // They should all be discarded since the host should be marked as down + // due to the connection returned with a network error. + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), 0); +} + +/** + * Verify that a connection returned with an error that does _not_ indicate + * the remote is unavailable does _not_ drop current generation connections to that remote. + */ +TEST_F(ConnectionPoolTest, OtherErrorsDontDropConns) { + auto pool = makePool(); + + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), 0U); + + constexpr size_t kSize = 100; + std::vector<ConnectionPool::ConnectionHandle> connections; + + // Ensure that no matter how we leave the test, we mark any + // checked out connections as OK before implicity returning them + // to the pool by destroying the 'connections' vector. Otherwise, + // this test would cause an invariant failure instead of a normal + // test failure if it fails, which would be confusing. + auto drainConnPool = [&] { + while (!connections.empty()) { + try { + ConnectionPool::ConnectionHandle conn = std::move(connections.back()); + connections.pop_back(); + doneWith(conn); + } catch (...) { + } + } + }; + const ScopeGuard guard(drainConnPool); + + auto now = Date_t::now(); + PoolImpl::setNow(now); + + auto checkOutConnections = [&] { + std::vector<unittest::ThreadAssertionMonitor> monitors(kSize); + for (size_t i = 0; i != kSize; ++i) { + ConnectionImpl::pushSetup(Status::OK()); + auto cb = [&](StatusWith<ConnectionPool::ConnectionHandle> swConn) { + monitors[i].exec([&]() { + ASSERT(swConn.isOK()); + connections.push_back(std::move(swConn.getValue())); + monitors[i].notifyDone(); + }); + }; + auto timeout = Milliseconds(5000); + pool->get_forTest(HostAndPort(), timeout, cb); + } + + for (auto& monitor : monitors) { + monitor.wait(); + } + + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), kSize); + }; + + // All three types of error that shouldn't result in us dropping connections - a non-network + // error; a network timeout error, and a network error that we can isolate to a specific + // connection. + std::array<ErrorCodes::Error, 3> errors = { + ErrorCodes::InternalError, ErrorCodes::NetworkTimeout, ErrorCodes::ConnectionError}; + for (size_t i = 0; i < errors.size(); ++i) { + // Check out kSize connections from the pool. + checkOutConnections(); + // Return one connection with a non-network error. + ConnectionPool::ConnectionHandle conn = std::move(connections.back()); + connections.pop_back(); + doneWithError(conn, {errors[i], "error"}); + + // We should still have all of the connections open, minus the one we just returned with an + // error. + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), kSize - 1); + + // Put the remaining connections back. + drainConnPool(); + + // They should all still be open. + ASSERT_EQ(pool->getNumConnectionsPerHost(HostAndPort()), kSize - 1); + } +} + +/** * Verify that providing different host and ports gives you different * connections. */ diff --git a/src/mongo/executor/connection_pool_tl.cpp b/src/mongo/executor/connection_pool_tl.cpp index e6917ff4342..49bccf97b38 100644 --- a/src/mongo/executor/connection_pool_tl.cpp +++ b/src/mongo/executor/connection_pool_tl.cpp @@ -33,6 +33,7 @@ #include "mongo/executor/connection_pool_tl.h" +#include "mongo/base/error_codes.h" #include "mongo/client/authenticate.h" #include "mongo/config.h" #include "mongo/db/auth/authorization_manager.h" @@ -171,7 +172,7 @@ public: explicit TLConnectionSetupHook(executor::NetworkConnectionHook* hookToWrap, bool x509AuthOnly) : _wrappedHook(hookToWrap), _x509AuthOnly(x509AuthOnly) {} - BSONObj augmentIsMasterRequest(const HostAndPort& remoteHost, BSONObj cmdObj) override { + BSONObj augmentHelloRequest(const HostAndPort& remoteHost, BSONObj cmdObj) override { BSONObjBuilder bob(std::move(cmdObj)); bob.append("hangUpOnStepDown", false); auto systemUser = internalSecurity.getUser(); @@ -189,9 +190,9 @@ public: } Status validateHost(const HostAndPort& remoteHost, - const BSONObj& isMasterRequest, - const RemoteCommandResponse& isMasterReply) override try { - const auto& reply = isMasterReply.data; + const BSONObj& helloRequest, + const RemoteCommandResponse& helloReply) override try { + const auto& reply = helloReply.data; // X.509 auth only means we only want to use a single mechanism regards of what hello says if (_x509AuthOnly) { @@ -215,7 +216,7 @@ public: if (!_wrappedHook) { return Status::OK(); } else { - return _wrappedHook->validateHost(remoteHost, isMasterRequest, isMasterReply); + return _wrappedHook->validateHost(remoteHost, helloRequest, helloReply); } } catch (const DBException& e) { return e.toStatus(); @@ -326,37 +327,42 @@ void TLConnection::setup(Milliseconds timeout, SetupCallback cb, std::string ins #endif // For transient connections, only use X.509 auth. - auto isMasterHook = std::make_shared<TLConnectionSetupHook>(_onConnectHook, x509AuthOnly); + auto helloHook = std::make_shared<TLConnectionSetupHook>(_onConnectHook, x509AuthOnly); AsyncDBClient::connect( _peer, _sslMode, _serviceContext, _reactor, timeout, _transientSSLContext) .thenRunOn(_reactor) .onError([](StatusWith<AsyncDBClient::Handle> swc) -> StatusWith<AsyncDBClient::Handle> { - return Status(ErrorCodes::HostUnreachable, swc.getStatus().reason()); + if (const Status& status = swc.getStatus(); + status.code() == ErrorCodes::ConnectionError) { + return status; + } else { + return Status(ErrorCodes::HostUnreachable, status.reason()); + } }) - .then([this, isMasterHook, instanceName = std::move(instanceName)]( + .then([this, helloHook, instanceName = std::move(instanceName)]( AsyncDBClient::Handle client) { _client = std::move(client); - return _client->initWireVersion(instanceName, isMasterHook.get()); + return _client->initWireVersion(instanceName, helloHook.get()); }) - .then([this, isMasterHook]() -> Future<bool> { + .then([this, helloHook]() -> Future<bool> { if (_skipAuth) { return false; } - return _client->completeSpeculativeAuth(isMasterHook->getSession(), + return _client->completeSpeculativeAuth(helloHook->getSession(), auth::getInternalAuthDB(), - isMasterHook->getSpeculativeAuthenticateReply(), - isMasterHook->getSpeculativeAuthType()); + helloHook->getSpeculativeAuthenticateReply(), + helloHook->getSpeculativeAuthType()); }) - .then([this, isMasterHook, authParametersProvider](bool authenticatedDuringConnect) { + .then([this, helloHook, authParametersProvider](bool authenticatedDuringConnect) { if (_skipAuth || authenticatedDuringConnect) { return Future<void>::makeReady(); } boost::optional<std::string> mechanism; - if (!isMasterHook->saslMechsForInternalAuth().empty()) - mechanism = isMasterHook->saslMechsForInternalAuth().front(); + if (!helloHook->saslMechsForInternalAuth().empty()) + mechanism = helloHook->saslMechsForInternalAuth().front(); return _client->authenticateInternal(std::move(mechanism), authParametersProvider); }) .then([this] { @@ -414,8 +420,7 @@ void TLConnection::refresh(Milliseconds timeout, RefreshCallback cb) { }); _client - ->runCommandRequest( - {_peer, std::string("admin"), BSON("isMaster" << 1), BSONObj(), nullptr}) + ->runCommandRequest({_peer, std::string("admin"), BSON("hello" << 1), BSONObj(), nullptr}) .then([](executor::RemoteCommandResponse response) { return Future<void>::makeReady(response.status); }) diff --git a/src/mongo/executor/network_connection_hook.h b/src/mongo/executor/network_connection_hook.h index 04de5571cb6..8510344c85a 100644 --- a/src/mongo/executor/network_connection_hook.h +++ b/src/mongo/executor/network_connection_hook.h @@ -54,16 +54,16 @@ public: virtual ~NetworkConnectionHook() = default; /** - * Optionally augments the isMaster request sent while initializing the wire protocol. + * Optionally augments the "hello" request sent while initializing the wire protocol. * * By default this will just return the cmdObj passed in unaltered. */ - virtual BSONObj augmentIsMasterRequest(const HostAndPort& remoteHost, BSONObj cmdObj) { + virtual BSONObj augmentHelloRequest(const HostAndPort& remoteHost, BSONObj cmdObj) { return cmdObj; } /** - * Runs optional validation logic on an isMaster reply from a remote host. If a non-OK + * Runs optional validation logic on an "hello" reply from a remote host. If a non-OK * Status is returned, it will be propagated up to the completion handler for the command * that initiated the request that caused this connection to be created. This will * be called once for each connection that is created, even if a remote host with the @@ -77,8 +77,8 @@ public: * std::terminate. */ virtual Status validateHost(const HostAndPort& remoteHost, - const BSONObj& isMasterRequest, - const RemoteCommandResponse& isMasterReply) = 0; + const BSONObj& helloRequest, + const RemoteCommandResponse& helloReply) = 0; /** * Generates a command to run on the remote host immediately after connecting to it. diff --git a/src/mongo/executor/network_interface.h b/src/mongo/executor/network_interface.h index 799f473a12d..1867a87e71c 100644 --- a/src/mongo/executor/network_interface.h +++ b/src/mongo/executor/network_interface.h @@ -34,6 +34,7 @@ #include <string> #include "mongo/bson/bsonobjbuilder.h" +#include "mongo/client/async_client.h" #include "mongo/executor/task_executor.h" #include "mongo/transport/baton.h" #include "mongo/util/fail_point.h" @@ -252,6 +253,46 @@ public: Milliseconds timeout, Status status) = 0; + /** + * An RAII type that NetworkInterface uses to allow users to access a stream corresponding to a + * single network-level (i.e. TCP, UDS) connection on which to run commands directly. Generally + * speaking, users should use the NetworkInterface itself to run commands to take advantage of + * connection-pooling, hedging, and other features - but for special cases where users need to + * borrow their own network-stream for manual use, they can lease one through this type. + * LeasedStreams are minimal and do not offer automated health-management/automated refreshing + * while on lease - users are responsible for examining the health of the stream as-needed if + * they desire. Users are also responsible for reporting on the health of stream before the + * lease ends so that it can be subsequently re-used; see comments below for detail. + */ + class LeasedStream { + public: + virtual ~LeasedStream() = default; + + // AsyncDBClient provides the mongoRPC-API for running commands over this stream. This + // stream owns the AsyncDBClient and no outstanding networking should be scheduled on the + // client when it is destroyed. + virtual AsyncDBClient* getClient() = 0; + + // Indicates that the user is done with this leased stream, and no failures on it occured. + // Users MUST call either this function or indicateFailure before the LeasedStream is + // destroyed. + virtual void indicateSuccess() = 0; + // Indicates that the stream is unhealthy (i.e. the user received a network error indicating + // the stream failed). This prevents the stream from being reused. + virtual void indicateFailure(Status) = 0; + // Indicates that the stream has successfully performed networking over the stream. Updates + // metadata indicating the last healthy networking over the stream so that appropriate + // health-checks can be done after the lease ends. + virtual void indicateUsed() = 0; + }; + + /** + * Lease a stream from this NetworkInterface for manual use. + */ + virtual SemiFuture<std::unique_ptr<LeasedStream>> leaseStream(const HostAndPort& hostAndPort, + transport::ConnectSSLMode sslMode, + Milliseconds timeout) = 0; + protected: NetworkInterface(); }; diff --git a/src/mongo/executor/network_interface_integration_test.cpp b/src/mongo/executor/network_interface_integration_test.cpp index fdf1ac6f8ba..1722084ac3b 100644 --- a/src/mongo/executor/network_interface_integration_test.cpp +++ b/src/mongo/executor/network_interface_integration_test.cpp @@ -50,6 +50,7 @@ #include "mongo/unittest/integration_test.h" #include "mongo/unittest/unittest.h" #include "mongo/util/assert_util.h" +#include "mongo/util/fail_point.h" #include "mongo/util/scopeguard.h" namespace mongo { @@ -160,7 +161,7 @@ public: } void setUp() override { - startNet(std::make_unique<WaitForIsMasterHook>(this)); + startNet(std::make_unique<WaitForHelloHook>(this)); } // NetworkInterfaceIntegrationFixture::tearDown() shuts down the NetworkInterface. We always @@ -253,33 +254,33 @@ public: return ++numCurrentOpRan; } - struct IsMasterData { + struct HelloData { BSONObj request; RemoteCommandResponse response; }; - IsMasterData waitForIsMaster() { + HelloData waitForHello() { stdx::unique_lock<Latch> lk(_mutex); - _isMasterCond.wait(lk, [this] { return _isMasterResult != boost::none; }); + _helloCondVar.wait(lk, [this] { return _helloResult != boost::none; }); - return std::move(*_isMasterResult); + return std::move(*_helloResult); } - bool hasIsMaster() { + bool hasHelloResult() { stdx::lock_guard<Latch> lk(_mutex); - return _isMasterResult != boost::none; + return _helloResult != boost::none; } private: - class WaitForIsMasterHook : public NetworkConnectionHook { + class WaitForHelloHook : public NetworkConnectionHook { public: - explicit WaitForIsMasterHook(NetworkInterfaceTest* parent) : _parent(parent) {} + explicit WaitForHelloHook(NetworkInterfaceTest* parent) : _parent(parent) {} Status validateHost(const HostAndPort& host, const BSONObj& request, - const RemoteCommandResponse& isMasterReply) override { + const RemoteCommandResponse& helloReply) override { stdx::lock_guard<Latch> lk(_parent->_mutex); - _parent->_isMasterResult = IsMasterData{request, isMasterReply}; - _parent->_isMasterCond.notify_all(); + _parent->_helloResult = HelloData{request, helloReply}; + _parent->_helloCondVar.notify_all(); return Status::OK(); } @@ -296,8 +297,8 @@ private: }; Mutex _mutex = MONGO_MAKE_LATCH("NetworkInterfaceTest::_mutex"); - stdx::condition_variable _isMasterCond; - boost::optional<IsMasterData> _isMasterResult; + stdx::condition_variable _helloCondVar; + boost::optional<HelloData> _helloResult; }; class NetworkInterfaceInternalClientTest : public NetworkInterfaceTest { @@ -328,7 +329,7 @@ TEST_F(NetworkInterfaceTest, CancelLocally) { auto deferred = runCommand(cbh, makeTestCommand(kMaxWait, makeEchoCmdObj())); - waitForIsMaster(); + waitForHello(); fpb->waitForTimesEntered(fpb.initialTimesEntered() + 1); @@ -503,13 +504,35 @@ TEST_F(NetworkInterfaceTest, LateCancel) { assertNumOps(0u, 0u, 0u, 1u); } +TEST_F(NetworkInterfaceTest, ConnectionErrorDropsSingleConnection) { + FailPoint* failPoint = + globalFailPointRegistry().find("transportLayerASIOasyncConnectReturnsConnectionError"); + auto timesEntered = failPoint->setMode(FailPoint::nTimes, 1); + + auto cbh = makeCallbackHandle(); + auto deferred = runCommand(cbh, makeTestCommand(kMaxWait, makeEchoCmdObj())); + // Wait for one of the connection attempts to fail with a `ConnectionError`. + failPoint->waitForTimesEntered(timesEntered + 1); + auto result = deferred.get(); + + ASSERT_OK(result.status); + ConnectionPoolStats stats; + net().appendConnectionStats(&stats); + + ASSERT_EQ(stats.totalCreated, 2); + ASSERT_EQ(stats.totalInUse + stats.totalAvailable + stats.totalRefreshing, 1); + // Connection dropped during finishRefresh, so the dropped connection still + // counts toward the refreshed counter. + ASSERT_EQ(stats.totalRefreshed, 2); +} + TEST_F(NetworkInterfaceTest, AsyncOpTimeout) { // Kick off operation auto cb = makeCallbackHandle(); auto request = makeTestCommand(Milliseconds{1000}, makeSleepCmdObj()); auto deferred = runCommand(cb, request); - waitForIsMaster(); + waitForHello(); auto result = deferred.get(); @@ -533,13 +556,19 @@ TEST_F(NetworkInterfaceTest, AsyncOpTimeoutWithOpCtxDeadlineSooner) { serviceContext->registerClientObserver(std::make_unique<LockerNoopClientObserver>()); auto client = serviceContext->makeClient("NetworkClient"); auto opCtx = client->makeOperationContext(); - opCtx->setDeadlineAfterNowBy(opCtxDeadline, ErrorCodes::ExceededTimeLimit); + + auto stopWatch = serviceContext->getPreciseClockSource()->makeStopWatch(); + opCtx->setDeadlineByDate(stopWatch.start() + opCtxDeadline, ErrorCodes::ExceededTimeLimit); auto request = makeTestCommand(requestTimeout, makeSleepCmdObj(), opCtx.get()); auto deferred = runCommand(cb, request); + // The time returned in result.elapsed is measured from when the command started, which happens + // in runCommand. The delay between setting the deadline on opCtx and starting the command can + // be long enough that the assertion about opCtxDeadline fails. + auto networkStartCommandDelay = stopWatch.elapsed(); - waitForIsMaster(); + waitForHello(); auto result = deferred.get(); @@ -551,9 +580,10 @@ TEST_F(NetworkInterfaceTest, AsyncOpTimeoutWithOpCtxDeadlineSooner) { ASSERT_EQ(ErrorCodes::ExceededTimeLimit, result.status); ASSERT(result.elapsed); + // check that the request timeout uses the smaller of the operation context deadline and // the timeout specified in the request constructor. - ASSERT_GTE(result.elapsed.value(), opCtxDeadline); + ASSERT_GTE(result.elapsed.value() + networkStartCommandDelay, opCtxDeadline); ASSERT_LT(result.elapsed.value(), requestTimeout); assertNumOps(0u, 1u, 0u, 0u); } @@ -569,12 +599,19 @@ TEST_F(NetworkInterfaceTest, AsyncOpTimeoutWithOpCtxDeadlineLater) { serviceContext->registerClientObserver(std::make_unique<LockerNoopClientObserver>()); auto client = serviceContext->makeClient("NetworkClient"); auto opCtx = client->makeOperationContext(); - opCtx->setDeadlineAfterNowBy(opCtxDeadline, ErrorCodes::ExceededTimeLimit); + + auto stopWatch = serviceContext->getPreciseClockSource()->makeStopWatch(); + opCtx->setDeadlineByDate(stopWatch.start() + opCtxDeadline, ErrorCodes::ExceededTimeLimit); + auto request = makeTestCommand(requestTimeout, makeSleepCmdObj(), opCtx.get()); auto deferred = runCommand(cb, request); + // The time returned in result.elapsed is measured from when the command started, which happens + // in runCommand. The delay between setting the deadline on opCtx and starting the command can + // be long enough that the assertion about opCtxDeadline fails. + auto networkStartCommandDelay = stopWatch.elapsed(); - waitForIsMaster(); + waitForHello(); auto result = deferred.get(); @@ -586,10 +623,12 @@ TEST_F(NetworkInterfaceTest, AsyncOpTimeoutWithOpCtxDeadlineLater) { ASSERT_EQ(ErrorCodes::NetworkInterfaceExceededTimeLimit, result.status); ASSERT(result.elapsed); + // check that the request timeout uses the smaller of the operation context deadline and // the timeout specified in the request constructor. ASSERT_GTE(duration_cast<Milliseconds>(result.elapsed.value()), requestTimeout); - ASSERT_LT(duration_cast<Milliseconds>(result.elapsed.value()), opCtxDeadline); + ASSERT_LT(duration_cast<Milliseconds>(result.elapsed.value() + networkStartCommandDelay), + opCtxDeadline); assertNumOps(0u, 1u, 0u, 0u); } @@ -733,13 +772,13 @@ TEST_F(NetworkInterfaceTest, SetAlarm) { } TEST_F(NetworkInterfaceInternalClientTest, - IsMasterRequestContainsOutgoingWireVersionInternalClientInfo) { + HelloRequestContainsOutgoingWireVersionInternalClientInfo) { auto deferred = runCommand(makeCallbackHandle(), makeTestCommand(kNoTimeout, makeEchoCmdObj())); - auto isMasterHandshake = waitForIsMaster(); + auto helloHandshake = waitForHello(); - // Verify that the isMaster reply has the expected internalClient data. + // Verify that the "hello" reply has the expected internalClient data. auto wireSpec = WireSpec::instance().get(); - auto internalClientElem = isMasterHandshake.request["internalClient"]; + auto internalClientElem = helloHandshake.request["internalClient"]; ASSERT_EQ(internalClientElem.type(), BSONType::Object); auto minWireVersionElem = internalClientElem.Obj()["minWireVersion"]; auto maxWireVersionElem = internalClientElem.Obj()["maxWireVersion"]; @@ -754,14 +793,14 @@ TEST_F(NetworkInterfaceInternalClientTest, assertNumOps(0u, 0u, 0u, 1u); } -TEST_F(NetworkInterfaceTest, IsMasterRequestMissingInternalClientInfoWhenNotInternalClient) { +TEST_F(NetworkInterfaceTest, HelloRequestMissingInternalClientInfoWhenNotInternalClient) { resetIsInternalClient(false); auto deferred = runCommand(makeCallbackHandle(), makeTestCommand(kNoTimeout, makeEchoCmdObj())); - auto isMasterHandshake = waitForIsMaster(); + auto helloHandshake = waitForHello(); - // Verify that the isMaster reply has the expected internalClient data. - ASSERT_FALSE(isMasterHandshake.request["internalClient"]); + // Verify that the "hello" reply has the expected internalClient data. + ASSERT_FALSE(helloHandshake.request["internalClient"]); // Verify that the ping op is counted as a success. auto res = deferred.get(); ASSERT(res.elapsed); @@ -939,6 +978,31 @@ TEST_F(NetworkInterfaceTest, TearDownWaitsForInProgress) { ASSERT_EQ(getInProgress(), 0); } +TEST_F(NetworkInterfaceTest, RunCommandOnLeasedStream) { + auto cs = fixture(); + auto target = cs.getServers().front(); + auto leasedStream = net().leaseStream(target, transport::kGlobalSSLMode, kNoTimeout).get(); + auto* client = leasedStream->getClient(); + + auto request = RemoteCommandRequest(target, "admin", makeEchoCmdObj(), nullptr, kNoTimeout); + auto deferred = client->runCommandRequest(request); + + auto res = deferred.get(); + + ASSERT(res.elapsed); + uassertStatusOK(res.status); + leasedStream->indicateSuccess(); + leasedStream->indicateUsed(); + + // This opmsg request expect the following reply, which is generated below + // { echo: { echo: 1, foo: "bar", $db: "admin" }, ok: 1.0 } + auto cmdObj = res.data.getObjectField("echo"); + ASSERT_EQ(1, cmdObj.getIntField("echo")); + ASSERT_EQ("bar"_sd, cmdObj.getStringField("foo")); + ASSERT_EQ("admin"_sd, cmdObj.getStringField("$db")); + ASSERT_EQ(1, res.data.getIntField("ok")); +} + } // namespace } // namespace executor } // namespace mongo diff --git a/src/mongo/executor/network_interface_mock.h b/src/mongo/executor/network_interface_mock.h index e096589d640..b8610d58d9d 100644 --- a/src/mongo/executor/network_interface_mock.h +++ b/src/mongo/executor/network_interface_mock.h @@ -41,6 +41,7 @@ #include "mongo/stdx/condition_variable.h" #include "mongo/stdx/unordered_map.h" #include "mongo/stdx/unordered_set.h" +#include "mongo/transport/mock_session.h" #include "mongo/util/clock_source.h" #include "mongo/util/clock_source_mock.h" #include "mongo/util/time_support.h" @@ -148,6 +149,20 @@ public: void testEgress(const HostAndPort&, transport::ConnectSSLMode, Milliseconds, Status) override {} + using LeasedStreamMaker = + std::function<std::unique_ptr<NetworkInterface::LeasedStream>(HostAndPort)>; + + void setLeasedStreamMaker(LeasedStreamMaker lsm) { + _leasedStreamMaker = std::move(lsm); + } + + SemiFuture<std::unique_ptr<NetworkInterface::LeasedStream>> leaseStream( + const HostAndPort& hp, transport::ConnectSSLMode, Milliseconds) override { + invariant(_leasedStreamMaker, + "Tried to lease a stream from NetworkInterfaceMock without providing one"); + return (*_leasedStreamMaker)(hp); + } + //////////////////////////////////////////////////////////////////////////////// // // Methods for simulating network operations and the passage of time. @@ -273,9 +288,9 @@ public: void runReadyNetworkOperations(); /** - * Sets the reply of the 'isMaster' handshake for a specific host. This reply will only + * Sets the reply of the 'hello' handshake for a specific host. This reply will only * be given to the 'validateHost' method of the ConnectionHook set on this object - NOT - * to the completion handlers of any 'isMaster' commands scheduled with 'startCommand'. + * to the completion handlers of any 'hello' commands scheduled with 'startCommand'. * * This reply will persist until it is changed again using this method. * @@ -446,6 +461,8 @@ private: // The handshake replies set for each host. stdx::unordered_map<HostAndPort, RemoteCommandResponse> _handshakeReplies; // (M) + + boost::optional<LeasedStreamMaker> _leasedStreamMaker; }; /** diff --git a/src/mongo/executor/network_interface_mock_test.cpp b/src/mongo/executor/network_interface_mock_test.cpp index c3f419a391f..7ab98d5a969 100644 --- a/src/mongo/executor/network_interface_mock_test.cpp +++ b/src/mongo/executor/network_interface_mock_test.cpp @@ -70,23 +70,23 @@ TEST_F(NetworkInterfaceMockTest, ConnectionHook) { Milliseconds(30)}; // need to copy as it will be moved - auto isMasterReplyData = BSON("iamyour" - << "father"); + auto helloReplyData = BSON("iamyour" + << "father"); - RemoteCommandResponse isMasterReply{isMasterReplyData.copy(), Milliseconds(20)}; + RemoteCommandResponse helloReply{helloReplyData.copy(), Milliseconds(20)}; - net().setHandshakeReplyForHost(testHost(), std::move(isMasterReply)); + net().setHandshakeReplyForHost(testHost(), std::move(helloReply)); // Since the contract of these methods is that they do not throw, we run the ASSERTs in // the test scope. net().setConnectionHook(makeTestHook( [&](const HostAndPort& remoteHost, const BSONObj&, - const RemoteCommandResponse& isMasterReply) { + const RemoteCommandResponse& helloReply) { validateCalled = true; hostCorrectForValidate = (remoteHost == testHost()); - replyCorrectForValidate = SimpleBSONObjComparator::kInstance.evaluate( - isMasterReply.data == isMasterReplyData); + replyCorrectForValidate = + SimpleBSONObjComparator::kInstance.evaluate(helloReply.data == helloReplyData); return Status::OK(); }, [&](const HostAndPort& remoteHost) { @@ -169,9 +169,8 @@ TEST_F(NetworkInterfaceMockTest, ConnectionHook) { TEST_F(NetworkInterfaceMockTest, ConnectionHookFailedValidation) { net().setConnectionHook(makeTestHook( - [&](const HostAndPort& remoteHost, - const BSONObj&, - const RemoteCommandResponse& isMasterReply) -> Status { + [&](const HostAndPort& remoteHost, const BSONObj&, const RemoteCommandResponse& helloReply) + -> Status { // We just need some obscure non-OK code. return {ErrorCodes::ConflictingOperationInProgress, "blah"}; }, @@ -199,7 +198,7 @@ TEST_F(NetworkInterfaceMockTest, ConnectionHookFailedValidation) { { net().enterNetwork(); // We should have short-circuited the network and immediately called the callback. - // If we change isMaster replies to go through the normal network mechanism, + // If we change "hello" replies to go through the normal network mechanism, // this test will need to change. ASSERT(!net().hasReadyRequests()); net().exitNetwork(); @@ -212,9 +211,8 @@ TEST_F(NetworkInterfaceMockTest, ConnectionHookFailedValidation) { TEST_F(NetworkInterfaceMockTest, ConnectionHookNoRequest) { bool makeRequestCalled = false; net().setConnectionHook(makeTestHook( - [&](const HostAndPort& remoteHost, - const BSONObj&, - const RemoteCommandResponse& isMasterReply) -> Status { return Status::OK(); }, + [&](const HostAndPort& remoteHost, const BSONObj&, const RemoteCommandResponse& helloReply) + -> Status { return Status::OK(); }, [&](const HostAndPort& remoteHost) -> StatusWith<boost::optional<RemoteCommandRequest>> { makeRequestCalled = true; return {boost::none}; @@ -248,9 +246,8 @@ TEST_F(NetworkInterfaceMockTest, ConnectionHookNoRequest) { TEST_F(NetworkInterfaceMockTest, ConnectionHookMakeRequestFails) { bool makeRequestCalled = false; net().setConnectionHook(makeTestHook( - [&](const HostAndPort& remoteHost, - const BSONObj&, - const RemoteCommandResponse& isMasterReply) -> Status { return Status::OK(); }, + [&](const HostAndPort& remoteHost, const BSONObj&, const RemoteCommandResponse& helloReply) + -> Status { return Status::OK(); }, [&](const HostAndPort& remoteHost) -> StatusWith<boost::optional<RemoteCommandRequest>> { makeRequestCalled = true; return {ErrorCodes::InvalidSyncSource, "blah"}; @@ -285,9 +282,8 @@ TEST_F(NetworkInterfaceMockTest, ConnectionHookMakeRequestFails) { TEST_F(NetworkInterfaceMockTest, ConnectionHookHandleReplyFails) { bool handleReplyCalled = false; net().setConnectionHook(makeTestHook( - [&](const HostAndPort& remoteHost, - const BSONObj&, - const RemoteCommandResponse& isMasterReply) -> Status { return Status::OK(); }, + [&](const HostAndPort& remoteHost, const BSONObj&, const RemoteCommandResponse& helloReply) + -> Status { return Status::OK(); }, [&](const HostAndPort& remoteHost) -> StatusWith<boost::optional<RemoteCommandRequest>> { return boost::make_optional<RemoteCommandRequest>({}); }, diff --git a/src/mongo/executor/network_interface_tl.cpp b/src/mongo/executor/network_interface_tl.cpp index 1b74e1e9ee5..0f23b2f61d7 100644 --- a/src/mongo/executor/network_interface_tl.cpp +++ b/src/mongo/executor/network_interface_tl.cpp @@ -33,12 +33,14 @@ #include <fmt/format.h> +#include "mongo/base/checked_cast.h" #include "mongo/config.h" #include "mongo/db/auth/security_token.h" #include "mongo/db/server_options.h" #include "mongo/db/wire_version.h" #include "mongo/executor/connection_pool_tl.h" #include "mongo/executor/hedging_metrics.h" +#include "mongo/executor/network_interface.h" #include "mongo/executor/network_interface_tl_gen.h" #include "mongo/logv2/log.h" #include "mongo/rpc/get_status_from_command_result.h" @@ -1349,5 +1351,32 @@ void NetworkInterfaceTL::dropConnections(const HostAndPort& hostAndPort) { _pool->dropConnections(hostAndPort); } +AsyncDBClient* NetworkInterfaceTL::LeasedStream::getClient() { + return checked_cast<connection_pool_tl::TLConnection*>(_conn.get())->client(); +} + +void NetworkInterfaceTL::LeasedStream::indicateSuccess() { + return _conn->indicateSuccess(); +} + +void NetworkInterfaceTL::LeasedStream::indicateFailure(Status status) { + _conn->indicateFailure(status); +} + +void NetworkInterfaceTL::LeasedStream::indicateUsed() { + _conn->indicateUsed(); +} + +SemiFuture<std::unique_ptr<NetworkInterface::LeasedStream>> NetworkInterfaceTL::leaseStream( + const HostAndPort& hostAndPort, transport::ConnectSSLMode sslMode, Milliseconds timeout) { + + return _pool->lease(hostAndPort, sslMode, timeout) + .thenRunOn(_reactor) + .then([](auto conn) -> std::unique_ptr<NetworkInterface::LeasedStream> { + auto ptr = std::make_unique<NetworkInterfaceTL::LeasedStream>(std::move(conn)); + return ptr; + }) + .semi(); +} } // namespace executor } // namespace mongo diff --git a/src/mongo/executor/network_interface_tl.h b/src/mongo/executor/network_interface_tl.h index 158e4b0a679..68c51d57aa7 100644 --- a/src/mongo/executor/network_interface_tl.h +++ b/src/mongo/executor/network_interface_tl.h @@ -31,9 +31,12 @@ #include <deque> +#include <boost/optional.hpp> + #include "mongo/client/async_client.h" #include "mongo/db/service_context.h" #include "mongo/executor/connection_pool.h" +#include "mongo/executor/connection_pool_tl.h" #include "mongo/executor/network_interface.h" #include "mongo/logv2/log_severity.h" #include "mongo/platform/mutex.h" @@ -103,6 +106,33 @@ public: Milliseconds timeout, Status status) override; + /** + * NetworkInterfaceTL's implementation of a leased network-stream + * provided for manual use outside of the NITL's usual RPC API. + * When this type is destroyed, the destructor of the ConnectionHandle + * member will return the connection to this NetworkInterface's ConnectionPool. + */ + class LeasedStream : public NetworkInterface::LeasedStream { + public: + AsyncDBClient* getClient() override; + + LeasedStream(ConnectionPool::ConnectionHandle&& conn) : _conn{std::move(conn)} {} + + // These pass-through indications of the health of the leased + // stream to the underlying ConnectionHandle + void indicateSuccess() override; + void indicateUsed() override; + void indicateFailure(Status) override; + + private: + ConnectionPool::ConnectionHandle _conn; + }; + + SemiFuture<std::unique_ptr<NetworkInterface::LeasedStream>> leaseStream( + const HostAndPort& hostAndPort, + transport::ConnectSSLMode sslMode, + Milliseconds timeout) override; + private: struct RequestState; struct RequestManager; @@ -342,8 +372,10 @@ private: std::unique_ptr<transport::TransportLayer> _ownedTransportLayer; transport::ReactorHandle _reactor; - mutable Mutex _mutex = - MONGO_MAKE_LATCH(HierarchicalAcquisitionLevel(3), "NetworkInterfaceTL::_mutex"); + // TODO SERVER-75830: This Mutex used to be at hierarcichal acquisition level 3. We temporary + // removed the level because it is sometimes acquired as part of task-scheduling when + // lower-level mutexes (like the ConnectionPool's) are held. + mutable Mutex _mutex = MONGO_MAKE_LATCH("NetworkInterfaceTL::_mutex"); const ConnectionPool::Options _connPoolOpts; std::unique_ptr<NetworkConnectionHook> _onConnectHook; std::shared_ptr<ConnectionPool> _pool; diff --git a/src/mongo/executor/pinned_connection_task_executor.cpp b/src/mongo/executor/pinned_connection_task_executor.cpp new file mode 100644 index 00000000000..2f90d464530 --- /dev/null +++ b/src/mongo/executor/pinned_connection_task_executor.cpp @@ -0,0 +1,443 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 "pinned_connection_task_executor.h" +#include "mongo/executor/network_interface.h" +#include "mongo/executor/thread_pool_task_executor.h" +#include "mongo/util/scoped_unlock.h" + +namespace mongo::executor { +/** + * Used as the state for callbacks _only_ for RPCs scheduled through this executor. + */ +class PinnedConnectionTaskExecutor::CallbackState : public TaskExecutor::CallbackState { + CallbackState(const CallbackState&) = delete; + CallbackState& operator=(const CallbackState&) = delete; + +public: + static std::shared_ptr<CallbackState> make(const RemoteCommandOnAnyCallbackFn& cb, + const BatonHandle& baton) { + return std::make_shared<CallbackState>(cb, baton); + } + + /** + * Do not call directly. Use make. + */ + CallbackState(const RemoteCommandOnAnyCallbackFn& cb, const BatonHandle& baton) + : callback(cb), baton(baton) {} + + virtual ~CallbackState() = default; + + bool isCanceled() const override { + MONGO_UNREACHABLE; + } + + void cancel() override { + MONGO_UNREACHABLE; + } + + void waitForCompletion() override { + MONGO_UNREACHABLE; + } + + // Run callback with a CallbackCanceled error. + static void runCallbackCanceled(stdx::unique_lock<Latch>& lk, + RequestAndCallback rcb, + TaskExecutor* exec) { + CallbackHandle cbHandle; + setCallbackForHandle(&cbHandle, rcb.second); + auto errorResponse = RemoteCommandOnAnyResponse(boost::none, kCallbackCanceledErrorStatus); + TaskExecutor::RemoteCommandOnAnyCallbackFn callback; + using std::swap; + swap(rcb.second->callback, callback); + ScopedUnlock guard(lk); + callback({exec, cbHandle, rcb.first, errorResponse}); + } + + // Run callback with the provided result. + static void runCallbackFinished(stdx::unique_lock<Latch>& lk, + RequestAndCallback rcb, + TaskExecutor* exec, + const StatusWith<RemoteCommandResponse>& result, + boost::optional<HostAndPort> targetUsed) { + // Convert the result into a RemoteCommandResponse unconditionally. + RemoteCommandResponse asRcr = + result.isOK() ? result.getValue() : RemoteCommandResponse(result.getStatus()); + // Convert the response into an OnAnyResponse using the provided target. + RemoteCommandOnAnyResponse asOnAnyRcr(targetUsed, asRcr); + CallbackHandle cbHandle; + setCallbackForHandle(&cbHandle, rcb.second); + TaskExecutor::RemoteCommandOnAnyCallbackFn callback; + using std::swap; + swap(rcb.second->callback, callback); + ScopedUnlock guard(lk); + callback({exec, cbHandle, rcb.first, asOnAnyRcr}); + } + + // All fields except for "canceled" are guarded by the owning task executor's _mutex. + enum class State { kWaiting, kRunning, kDone, kCanceled }; + + RemoteCommandOnAnyCallbackFn callback; + boost::optional<stdx::condition_variable> finishedCondition; + State state{State::kWaiting}; + bool isNetworkOperation = true; + bool startedNetworking = false; + BatonHandle baton; +}; + +PinnedConnectionTaskExecutor::PinnedConnectionTaskExecutor( + const std::shared_ptr<TaskExecutor>& executor, NetworkInterface* net) + : _executor(executor), _net(net), _cancellationExecutor(executor) {} + +PinnedConnectionTaskExecutor::~PinnedConnectionTaskExecutor() { + shutdown(); + join(); +} + +Date_t PinnedConnectionTaskExecutor::now() { + return _executor->now(); +} + +StatusWith<TaskExecutor::EventHandle> PinnedConnectionTaskExecutor::makeEvent() { + return _executor->makeEvent(); +} + +void PinnedConnectionTaskExecutor::signalEvent(const EventHandle& event) { + return _executor->signalEvent(event); +} + +StatusWith<TaskExecutor::CallbackHandle> PinnedConnectionTaskExecutor::onEvent( + const EventHandle& event, CallbackFn&& work) { + return _executor->onEvent(event, std::move(work)); +} + +void PinnedConnectionTaskExecutor::waitForEvent(const EventHandle& event) { + _executor->waitForEvent(event); +} + +StatusWith<stdx::cv_status> PinnedConnectionTaskExecutor::waitForEvent(OperationContext* opCtx, + const EventHandle& event, + Date_t deadline) { + return _executor->waitForEvent(opCtx, event, deadline); +} + +StatusWith<TaskExecutor::CallbackHandle> PinnedConnectionTaskExecutor::scheduleWork( + CallbackFn&& work) { + return _executor->scheduleWork(std::move(work)); +} + +StatusWith<TaskExecutor::CallbackHandle> PinnedConnectionTaskExecutor::scheduleWorkAt( + Date_t when, CallbackFn&& work) { + return _executor->scheduleWorkAt(when, std::move(work)); +} + +StatusWith<TaskExecutor::CallbackHandle> PinnedConnectionTaskExecutor::scheduleRemoteCommandOnAny( + const RemoteCommandRequestOnAny& requestOnAny, + const RemoteCommandOnAnyCallbackFn& cb, + const BatonHandle& baton) { + + stdx::unique_lock<Latch> lk{_mutex}; + if (_state != State::running) { + return {ErrorCodes::ShutdownInProgress, "Shutdown in progress"}; + } + invariant(requestOnAny.target.size() == 1, + "RPCs scheduled through PinnedConnectionTaskExecutor can only target a single host."); + RemoteCommandRequest req = RemoteCommandRequest(requestOnAny, 0); + auto state = PinnedConnectionTaskExecutor::CallbackState::make(cb, baton); + _requestQueue.push_back({req, state}); + + CallbackHandle cbHandle; + setCallbackForHandle(&cbHandle, state); + + if (!_isDoingNetworking) { + _doNetworking(std::move(lk)); + } + + return cbHandle; +} + +void PinnedConnectionTaskExecutor::_cancel(WithLock, CallbackState* cbState) { + switch (cbState->state) { + case CallbackState::State::kWaiting: + // Just set the state to canceled. The callback will be run with an + // error status once it reaches the front of the queue. + cbState->state = CallbackState::State::kCanceled; + break; + case CallbackState::State::kRunning: { + // Cancel the ongoing operation. + cbState->state = CallbackState::State::kCanceled; + if (_stream) { + auto client = _stream->getClient(); + client->cancel(cbState->baton); + } + break; + } + case CallbackState::State::kCanceled: + [[fallthrough]]; + case CallbackState::State::kDone: + // Nothing to do. + break; + } +} + +void PinnedConnectionTaskExecutor::cancel(const CallbackHandle& cbHandle) { + auto cbState = + dynamic_cast<PinnedConnectionTaskExecutor::CallbackState*>(getCallbackFromHandle(cbHandle)); + if (!cbState) { + // Defer to underlying for non-RPC. + _executor->cancel(cbHandle); + return; + } + stdx::lock_guard lk(_mutex); + return _cancel(std::move(lk), cbState); +} + +ExecutorFuture<void> PinnedConnectionTaskExecutor::_ensureStream( + WithLock, HostAndPort target, Milliseconds timeout, transport::ConnectSSLMode sslMode) { + if (!_stream) { + auto streamFuture = _net->leaseStream(target, sslMode, timeout); + // If the stream is ready, send the RPC immediately by continuing inline. + if (streamFuture.isReady()) { + auto stream = std::move(streamFuture).getNoThrow(); + if (!stream.isOK()) { + // Propogate the error down the future chain. + return ExecutorFuture<void>(*_executor, stream.getStatus()); + } + _stream = std::move(stream.getValue()); + return ExecutorFuture<void>(*_executor); + } + // Otherwise continue on the networking reactor once the stream is ready. + return std::move(streamFuture) + .thenRunOn(*_executor) + .then([this](std::unique_ptr<NetworkInterface::LeasedStream> stream) { + stdx::lock_guard lk{_mutex}; + _stream = std::move(stream); + }); + } + + auto remote = _stream->getClient()->remote(); + using namespace fmt::literals; + invariant( + target == remote, + "Attempted to schedule RPC to {} on TaskExecutor that had pinned connection to {}"_format( + target, remote)); + return ExecutorFuture<void>(*_executor); +} + +Future<executor::RemoteCommandResponse> PinnedConnectionTaskExecutor::_runSingleCommand( + RemoteCommandRequest command, std::shared_ptr<CallbackState> cbState) { + stdx::lock_guard lk{_mutex}; + if (auto& state = cbState->state; MONGO_unlikely(state == CallbackState::State::kCanceled)) { + // It's possible this callback was canceled after it was moved + // out of the queue, but before we actually started work on the client. + // In that case, don't run it. + return kCallbackCanceledErrorStatus; + } + auto client = _stream->getClient(); + cbState->startedNetworking = true; + return client->runCommandRequest(command, cbState->baton); +} + +boost::optional<PinnedConnectionTaskExecutor::RequestAndCallback> +PinnedConnectionTaskExecutor::_getFirstUncanceledRequest(stdx::unique_lock<Latch>& lk) { + while (!_requestQueue.empty()) { + auto req = std::move(_requestQueue.front()); + _requestQueue.pop_front(); + if (req.second->state == CallbackState::State::kCanceled) { + CallbackState::runCallbackCanceled(lk, req, this); + } else { + return req; + } + } + return boost::none; +} + +void PinnedConnectionTaskExecutor::_doNetworking(stdx::unique_lock<Latch>&& lk) { + _isDoingNetworking = true; + // Find the first non-canceled request. + boost::optional<RequestAndCallback> maybeReqToRun = _getFirstUncanceledRequest(lk); + if (!maybeReqToRun) { + // No non-canceled requests. Stop doing networking. + _isDoingNetworking = false; + invariant(_requestQueue.empty()); + _requestQueueEmptyCV.notify_all(); + return; + } + auto req = *maybeReqToRun; + // Set req state to running + invariant(req.second->state == CallbackState::State::kWaiting); + req.second->state = CallbackState::State::kRunning; + auto streamFut = _ensureStream(lk, req.first.target, req.first.timeout, req.first.sslMode); + // Stash the in-progress operation before releasing the lock so we can + // access it if we're shutdown while it's in-progress. + _inProgressRequest = req.second; + lk.unlock(); + std::move(streamFut) + .then([req, this]() { return _runSingleCommand(req.first, req.second); }) + .thenRunOn(makeGuaranteedExecutor(req.second->baton, _cancellationExecutor)) + .getAsync([req, this, self = shared_from_this()](StatusWith<RemoteCommandResponse> result) { + stdx::unique_lock<Latch> lk{_mutex}; + _inProgressRequest.reset(); + // If we used the _stream, update it accordingly. + if (req.second->startedNetworking) { + if (auto status = result.getStatus(); status.isOK()) { + _stream->indicateUsed(); + _stream->indicateSuccess(); + } else { + // We didn't get a response from the remote. + // We assume the stream is broken and therefore can do no more work. Notify the + // stream of the failure, destroy it, and shutdown. + _stream->indicateFailure(status); + _stream.reset(); + _shutdown(lk); + } + } + // Now run the completion callback for the command. + if (auto& state = req.second->state; + MONGO_unlikely(state == CallbackState::State::kCanceled)) { + CallbackState::runCallbackCanceled(lk, req, this); + } else { + invariant(state == CallbackState::State::kRunning); + // Three possibilities here: we either finished the RPC + // successfully, got a local error from the stream after + // attempting to start networking, or never were able to acquire a + // stream. In any case, we first complete the current request + // by invoking it's callback: + state = CallbackState::State::kDone; + // Get the target if we successfully acquired a stream. + boost::optional<HostAndPort> target = boost::none; + if (_stream) { + target = _stream->getClient()->remote(); + } + CallbackState::runCallbackFinished(lk, req, this, result, target); + } + // If we weren't able to acquire a stream, shut-down. + if (!_stream) { + _shutdown(lk); + } + _isDoingNetworking = false; + if (!_requestQueue.empty()) { + return _doNetworking(std::move(lk)); + } + _requestQueueEmptyCV.notify_all(); + }); +} + +void PinnedConnectionTaskExecutor::_shutdown(WithLock lk) { + if (_state != State::running) { + return; + } + _state = State::joinRequired; + _executor->shutdown(); + for (auto&& [_, cbState] : _requestQueue) { + _cancel(lk, cbState.get()); + } + if (_isDoingNetworking && _inProgressRequest) { + // Cancel the in-progress request that was already popped from the queue. + _cancel(lk, _inProgressRequest.get()); + } +} + +void PinnedConnectionTaskExecutor::shutdown() { + stdx::lock_guard lk(_mutex); + _shutdown(lk); +} + +// May be called by any thread that wishes to wait until this executor is done shutting down. +// Any thread that calls this will block until no work remains scheduled but not completed +// on this executor. After join() completes, the state if this executor will be 'shutdownComplete'. +void PinnedConnectionTaskExecutor::join() { + stdx::unique_lock lk(_mutex); + if (_state == State::shutdownComplete) { + return; + } + invariant(_state == State::joinRequired || _state == State::joining); + _state = State::joining; + + _requestQueueEmptyCV.wait(lk, + [this]() { return _requestQueue.empty() && !_isDoingNetworking; }); + + _executor->join(); + + _state = State::shutdownComplete; + return; +} + +SharedSemiFuture<void> PinnedConnectionTaskExecutor::joinAsync() { + MONGO_UNREACHABLE; +} + +bool PinnedConnectionTaskExecutor::isShuttingDown() const { + stdx::lock_guard lk(_mutex); + return _state != State::running; +} + + +// Below are the portions of the TaskExecutor API that are illegal to use through +// PinnedCursorTaskExecutor and/or are unimplemented at this time. +void PinnedConnectionTaskExecutor::wait(const CallbackHandle& cbHandle, + Interruptible* interruptible) { + MONGO_UNREACHABLE; +} + +StatusWith<TaskExecutor::CallbackHandle> +PinnedConnectionTaskExecutor::scheduleExhaustRemoteCommandOnAny( + const RemoteCommandRequestOnAny& request, + const RemoteCommandOnAnyCallbackFn& cb, + const BatonHandle& baton) { + MONGO_UNREACHABLE; +} + +bool PinnedConnectionTaskExecutor::hasTasks() { + stdx::lock_guard lk(_mutex); + return (!_requestQueue.empty()) || _executor->hasTasks(); +} + +void PinnedConnectionTaskExecutor::startup() { + MONGO_UNREACHABLE; +} + +void PinnedConnectionTaskExecutor::appendDiagnosticBSON(mongo::BSONObjBuilder* builder) const { + MONGO_UNREACHABLE; +} + + +void PinnedConnectionTaskExecutor::appendConnectionStats(ConnectionPoolStats* stats) const { + MONGO_UNREACHABLE; +} + +void PinnedConnectionTaskExecutor::dropConnections(const HostAndPort& hostAndPort) { + MONGO_UNREACHABLE; +} + +void PinnedConnectionTaskExecutor::appendNetworkInterfaceStats(BSONObjBuilder& bob) const { + MONGO_UNREACHABLE; +} + +} // namespace mongo::executor diff --git a/src/mongo/executor/pinned_connection_task_executor.h b/src/mongo/executor/pinned_connection_task_executor.h new file mode 100644 index 00000000000..9d423931327 --- /dev/null +++ b/src/mongo/executor/pinned_connection_task_executor.h @@ -0,0 +1,196 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 <memory> + +#include "mongo/executor/network_interface.h" +#include "mongo/executor/scoped_task_executor.h" +#include "mongo/executor/task_executor.h" + +namespace mongo::executor { + +/** + * Implementation of a TaskExecutor that provides the ability to schedule RPC/networking on the same + * underlying network connection. The PinnedTaskExecutor is constructed from another TaskExecutor, + * and uses that TaskExecutor's ThreadPool and NetworkInterface/networking reactor to perform work. + * Specifically: + * - Functions that schedule work or manage events that happen locally, without going over the + * network, are passed-through to the underlying TaskExecutor (i.e. scheduleWork, + * makeEvent, waitForEvent). + * - Functions that involve scheduling RPC/networking are all run on the same underlying + * network-connection (i.e. TCP/Unix Domain Socket). + * Note that this means that the PinnedConnectionTaskExecutor can only speak to one host over its + * entire lifetime! If you need to speak to a different host, you need a different connection, so + * construct a *new* PinnedCursorTaskExecutor from the underlying executor. + * + * Certain methods are illegal to call. startup() is illegal to call because the TaskExecutor + * passed to PinnedConnectionTaskExecutor should be started-up prior to this object's construction, + * and no additional startup is needed. + * Additionally, diagnostic and network management methods: + * - appendDiagnosticBSON() + * - appendConnectionStats() + * - dropConnections() + * - appendNetworkInterfaceStats() + * are illegal to call because this TaskExecutor provides a distinct networking API. Gather + * diagnostics from the underlying TaskExecutor instead if needed. + * + * This type uses ScopedTaskExecutor to proxy work to the underlying TaskExecutor it is + * constructed from. This means that shutdown() and join() address only tasks dispatched + * through this executor, rather than passing through to the underlying executor. + * + * Note! The executor that this PinnedConnectionTaskExecutor is constructed from _must_ + * out-life it - i.e. this PinnedConnectionTaskExecutor must be shutdown and joined + * before the underlying executor is. This is because this type must have access + * to the underlying thread pool to complete cancellation tasks as it shuts down. + * + * Exhaust commands are not supported at this time. + */ +class PinnedConnectionTaskExecutor final : public TaskExecutor { + PinnedConnectionTaskExecutor(const PinnedConnectionTaskExecutor&) = delete; + PinnedConnectionTaskExecutor& operator=(const PinnedConnectionTaskExecutor&) = delete; + +public: + // The provided NetworkInterface should be owned by the provided TaskExecutor, and + // must outlive this type. + PinnedConnectionTaskExecutor(const std::shared_ptr<TaskExecutor>& executor, + NetworkInterface* net); + + ~PinnedConnectionTaskExecutor(); + // Startup is illegal to call, as the provided executor should already be started-up. + void startup() override; + void shutdown() override; + void join() override; + SharedSemiFuture<void> joinAsync() override; + bool isShuttingDown() const override; + + // These pass-through to the underlying TaskExecutor. + Date_t now() override; + StatusWith<EventHandle> makeEvent() override; + void signalEvent(const EventHandle& event) override; + StatusWith<CallbackHandle> onEvent(const EventHandle& event, CallbackFn&& work) override; + void waitForEvent(const EventHandle& event) override; + StatusWith<stdx::cv_status> waitForEvent(OperationContext* opCtx, + const EventHandle& event, + Date_t deadline) override; + StatusWith<CallbackHandle> scheduleWork(CallbackFn&& work) override; + StatusWith<CallbackHandle> scheduleWorkAt(Date_t when, CallbackFn&& work) override; + + // This type provides special connection-pinning behavior for RPC functionality here. + StatusWith<CallbackHandle> scheduleRemoteCommandOnAny( + const RemoteCommandRequestOnAny& request, + const RemoteCommandOnAnyCallbackFn& cb, + const BatonHandle& baton = nullptr) override; + + StatusWith<CallbackHandle> scheduleExhaustRemoteCommandOnAny( + const RemoteCommandRequestOnAny& request, + const RemoteCommandOnAnyCallbackFn& cb, + const BatonHandle& baton = nullptr) override; + + // When cancel() is passed a CallbackHandle that was returned from schedule{Work}()/onEvent(), + // cancellation is passed-through to the underlying executor. If the CallbackHandle was returned + // from scheduleRemoteCommand then the executor will cancel the RPC attempt. + void cancel(const CallbackHandle& cbHandle) override; + + // Wait is unimplemented at this time. + void wait(const CallbackHandle& cbHandle, + Interruptible* interruptible = Interruptible::notInterruptible()) override; + + // Illegal to call because the view does not track it's portion of the underlying TaskExecutor's + // resources. + void appendConnectionStats(ConnectionPoolStats*) const override; + void appendNetworkInterfaceStats(BSONObjBuilder&) const override; + void appendDiagnosticBSON(BSONObjBuilder*) const override; + void dropConnections(const HostAndPort&) override; + bool hasTasks() override; + +private: + // Ensures _stream is initialized with a valid LeasedStream to `target`. + // If we already have a _stream when this function is called, ensures the + // remote is `target` and returns a ready-future. Otherwise asynchronously + // initailizes _stream and returns a future that resolves once _stream is ready. + ExecutorFuture<void> _ensureStream(WithLock, + HostAndPort target, + Milliseconds timeout, + transport::ConnectSSLMode sslMode); + + // Start processing pending/queued RPCs. + void _doNetworking(stdx::unique_lock<Latch>&&); + + // CallbackState for RPCs. Non-RPC callbacks use the CallbackState from the _underlyingExecutor. + class CallbackState; + + // Invoke the RPC and return a future of its response. + Future<RemoteCommandResponse> _runSingleCommand(RemoteCommandRequest command, + std::shared_ptr<CallbackState> cbState); + + void _shutdown(WithLock); + + // Alias for an RPC request and the associated CallbackState. + using RequestAndCallback = std::pair<RemoteCommandRequest, std::shared_ptr<CallbackState>>; + + // Helper to cancel a CallbackState from this executor. + void _cancel(WithLock, CallbackState*); + + // Helper that walks the _requestQueue in-order, completing any canceled callbacks, until + // it finds the first uncanceled one (if any), which it returns. + boost::optional<RequestAndCallback> _getFirstUncanceledRequest(stdx::unique_lock<Latch>&); + + // Synchronizes access to the _requestQueue, _stream, and _isDoingNetworking variables, as well + // as all CallbackState members. + mutable Mutex _mutex; + + ScopedTaskExecutor _executor; + // Owned by the TaskExecutor backing _executor above. Since ScopedTaskExecutor keeps a + // shared_ptr to it's backing TaskExecutor, _net will remain valid for at least the lifetime of + // _executor. + NetworkInterface* _net; + + // This is the same executor that the ScopedTaskExecutor above provides a view over. We keep + // a pointer to it so that we can run cancellation tasks even after the ScopedTaskExecutor + // is shut down. This should _only_ be used to guarantee cancellation tasks will run, even + // after shutdown is called on this type! + std::shared_ptr<TaskExecutor> _cancellationExecutor; + + // Queue of pending/not-yet-started RPC requests and corresponding completion callbacks + // scheduled on this executor. + std::deque<RequestAndCallback> _requestQueue; + stdx::condition_variable _requestQueueEmptyCV; + // Pinned-connection leased from _underlyingNet to run all RPCs through this executor. + // Initialized upon the execution of the first scheduled RPC, and subsequently re-used for all + // RPCs scheduled through this executor. + std::unique_ptr<NetworkInterface::LeasedStream> _stream; + bool _isDoingNetworking{false}; + std::shared_ptr<CallbackState> _inProgressRequest; + + enum class State { running, joinRequired, joining, shutdownComplete }; + State _state = State::running; +}; + +} // namespace mongo::executor diff --git a/src/mongo/executor/pinned_connection_task_executor_factory.cpp b/src/mongo/executor/pinned_connection_task_executor_factory.cpp new file mode 100644 index 00000000000..e5871adbaa0 --- /dev/null +++ b/src/mongo/executor/pinned_connection_task_executor_factory.cpp @@ -0,0 +1,54 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 <memory> + +#include "mongo/executor/pinned_connection_task_executor_factory.h" + +#include "mongo/executor/pinned_connection_task_executor.h" +#include "mongo/executor/thread_pool_task_executor.h" + +namespace mongo { +namespace executor { + +std::shared_ptr<TaskExecutor> makePinnedConnectionTaskExecutor(std::shared_ptr<TaskExecutor> exec, + NetworkInterface* net) { + return std::make_shared<PinnedConnectionTaskExecutor>(std::move(exec), net); +} + +std::shared_ptr<TaskExecutor> makePinnedConnectionTaskExecutor(std::shared_ptr<TaskExecutor> exec) { + auto tpte = dynamic_cast<ThreadPoolTaskExecutor*>(exec.get()); + invariant(tpte, + "Connection-pinning task executors can only be constructed from " + "ThreadPoolTaskExecutor unless an explicit NetworkInterface is provided."); + return makePinnedConnectionTaskExecutor(std::move(exec), tpte->_net.get()); +} + +} // namespace executor +} // namespace mongo diff --git a/src/mongo/executor/pinned_connection_task_executor_factory.h b/src/mongo/executor/pinned_connection_task_executor_factory.h new file mode 100644 index 00000000000..dbf79481bff --- /dev/null +++ b/src/mongo/executor/pinned_connection_task_executor_factory.h @@ -0,0 +1,56 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 <memory> + +#include "mongo/executor/network_interface.h" +#include "mongo/executor/task_executor.h" + +namespace mongo { +namespace executor { + +/** + * Returns a new TaskExecutor that does all of its RPC execution over the same transport session. + * The returned executor uses `exec`'s execution resources and acquires the transport session from + * `net`. + */ +std::shared_ptr<TaskExecutor> makePinnedConnectionTaskExecutor(std::shared_ptr<TaskExecutor> exec, + NetworkInterface* net); + +/** + * Returns a new TaskExecutor that does all of its RPC execution over the same transport session. + * The provided executor _must_ be a ThreadPoolTaskExecutor, and its underlying execution and + * network resources will be used by the returned executor. + */ +std::shared_ptr<TaskExecutor> makePinnedConnectionTaskExecutor(std::shared_ptr<TaskExecutor> exec); + +} // namespace executor +} // namespace mongo diff --git a/src/mongo/executor/pinned_connection_task_executor_test.cpp b/src/mongo/executor/pinned_connection_task_executor_test.cpp new file mode 100644 index 00000000000..6fa8956f5dd --- /dev/null +++ b/src/mongo/executor/pinned_connection_task_executor_test.cpp @@ -0,0 +1,460 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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 "pinned_connection_task_executor_test_fixture.h" + +#include "mongo/rpc/get_status_from_command_result.h" +#include "mongo/rpc/op_msg_rpc_impls.h" +#include "mongo/unittest/death_test.h" +#include "mongo/unittest/thread_assertion_monitor.h" +#include "mongo/unittest/unittest.h" +#include "mongo/util/assert_util.h" +#include "mongo/util/duration.h" +#include "mongo/util/future.h" +#include "mongo/util/future_impl.h" +#include "mongo/util/net/hostandport.h" +#include "mongo/util/time_support.h" + +namespace mongo::executor { +namespace { + +RemoteCommandRequest makeRCR(HostAndPort remote, BSONObj extraFields) { + return RemoteCommandRequest(remote, "admin", BSON("hello" << 1), extraFields, nullptr); +}; + +void assertMessageBodyCameFromRequest(Message m, RemoteCommandRequest rcr) { + auto opMsg = OpMsgRequest::parse(m); + auto expectedOpMsg = OpMsgRequest::fromDBAndBody( + std::move(rcr.dbname), std::move(rcr.cmdObj), std::move(rcr.metadata)); + ASSERT_BSONOBJ_EQ(opMsg.body, expectedOpMsg.body); +} + +void assertMessageBodyAndDBName(Message m, BSONObj body, BSONObj metadata, std::string dbName) { + auto opMsg = OpMsgRequest::parse(m); + auto expectedOpMsg = OpMsgRequest::fromDBAndBody(dbName, body, metadata); + ASSERT_BSONOBJ_EQ(opMsg.body, expectedOpMsg.body); +} + +Message makeOkReplyMessage() { + rpc::OpMsgReplyBuilder replyBuilder; + replyBuilder.setCommandReply(BSONObj()); + return replyBuilder.done(); +} + +Message makeErrorReplyMessage(Status error) { + rpc::OpMsgReplyBuilder replyBuilder; + replyBuilder.setCommandReply(error); + return replyBuilder.done(); +} + +TEST_F(PinnedConnectionTaskExecutorTest, RunSingleCommandOverSession) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + + auto rcr = makeRCR(remote, BSONObj()); + auto pf = makePromiseFuture<void>(); + + ASSERT_OK(pinnedTE + ->scheduleRemoteCommand(rcr, + [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + pf.promise.setWith( + [&] { return args.response.status; }); + }) + .getStatus()); + // We first expect sink message to be called and to see the hello + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyCameFromRequest(m, rcr); + return Status::OK(); + }); + // Now we expect source message to be called and provide the response + expectSourceMessage([&]() { + auto message = makeOkReplyMessage(); + message.header().setResponseToMsgId(responseToId); + return message; + }); + + ASSERT_OK(pf.future.getNoThrow()); + pinnedTE->shutdown(); + pinnedTE->join(); +} + +// Test we can schedule multiple RPC on the executor, and that they then +// run serially over the same transport session. +TEST_F(PinnedConnectionTaskExecutorTest, RunTwoRemoteCommandsSimultaneously) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + + // Schedule two RPCs + std::vector<Future<void>> results; + for (int i = 0; i < 2; ++i) { + auto promise = std::make_shared<Promise<void>>(NonNullPromiseTag{}); + results.push_back(promise->getFuture()); + auto extraFields = BSON("forTest" << i); + ASSERT_OK( + pinnedTE + ->scheduleRemoteCommand( + makeRCR(remote, extraFields), + [p = std::move(promise)](const TaskExecutor::RemoteCommandCallbackArgs& args) { + p->setWith([&] { return args.response.status; }); + }) + .getStatus()); + } + ASSERT_EQ(2, results.size()); + for (int i = 0; i < 2; ++i) { + auto pf = makePromiseFuture<void>(); + // We first expect sink message to be called and to see the i'th request + // (All i requests should appear on our same mocked session). + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyAndDBName(m, BSON("hello" << 1), BSON("forTest" << i), "admin"); + pf.promise.emplaceValue(); + return Status::OK(); + }); + pf.future.get(); + // Now we expect source message to be called and provide the response + expectSourceMessage([&]() { + auto message = makeOkReplyMessage(); + message.header().setResponseToMsgId(responseToId); + return message; + }); + // I'th command should be completed: + ASSERT_OK(results[i].getNoThrow()); + } + pinnedTE->shutdown(); + pinnedTE->join(); +} + +TEST_F(PinnedConnectionTaskExecutorTest, FailCommandRemotelyDoesntBreakOtherCommands) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + // + // Schedule two RPCs + std::vector<Future<BSONObj>> results; + for (int i = 0; i < 2; ++i) { + auto promise = std::make_shared<Promise<BSONObj>>(NonNullPromiseTag{}); + results.push_back(promise->getFuture()); + auto extraFields = BSON("forTest" << i); + ASSERT_OK( + pinnedTE + ->scheduleRemoteCommand( + makeRCR(remote, extraFields), + [p = std::move(promise)](const TaskExecutor::RemoteCommandCallbackArgs& args) { + if (args.response.isOK()) { + p->emplaceValue(args.response.data); + } else { + p->setError(args.response.status); + } + }) + .getStatus()); + } + ASSERT_EQ(2, results.size()); + + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyAndDBName(m, BSON("hello" << 1), BSON("forTest" << 0), "admin"); + return Status::OK(); + }); + // Fail the first request + Status testFailure{ErrorCodes::BadValue, "test failure"}; + expectSourceMessage([&]() { + auto message = makeErrorReplyMessage(testFailure); + message.header().setResponseToMsgId(responseToId); + return message; + }); + auto remoteErr = results[0].getNoThrow().getValue(); + ASSERT_EQ(getStatusFromCommandResult(remoteErr), testFailure); + + // Second command should still be able to succeed: + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyAndDBName(m, BSON("hello" << 1), BSON("forTest" << 1), "admin"); + return Status::OK(); + }); + expectSourceMessage([&]() { + auto message = makeOkReplyMessage(); + message.header().setResponseToMsgId(responseToId); + return message; + }); + auto success = results[1].getNoThrow().getValue(); + ASSERT_EQ(Status::OK(), getStatusFromCommandResult(success)); + + pinnedTE->shutdown(); + pinnedTE->join(); +} + +DEATH_TEST_REGEX_F( + PinnedConnectionTaskExecutorTest, + SchedulingCommandOnDifferentHostFails, + R"#(Attempted to schedule RPC to (\S+):(\d+) on TaskExecutor that had pinned connection to (\S+):(\d+))#") { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + HostAndPort otherRemote("otherHost"); + + // Schedule two RPCs + auto pf = makePromiseFuture<void>(); + ASSERT_OK(pinnedTE + ->scheduleRemoteCommand(makeRCR(remote, {}), + [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + pf.promise.setWith( + [&] { return args.response.status; }); + }) + .getStatus()); + auto pfTwo = makePromiseFuture<void>(); + ASSERT_OK(pinnedTE + ->scheduleRemoteCommand(makeRCR(otherRemote, {}), + [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + pfTwo.promise.setWith( + [&] { return args.response.status; }); + }) + .getStatus()); + // first command runs OK + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + return Status::OK(); + }); + expectSourceMessage([&]() { + auto reply = makeOkReplyMessage(); + reply.header().setResponseToMsgId(responseToId); + return reply; + }); + ASSERT_OK(pf.future.getNoThrow()); + + // Second command should invariant once the PCTE attempts to run it, because it has a different + // remote target. + // Should never be fulfilled. + ASSERT_OK(pfTwo.future.getNoThrow()); +} + +TEST_F(PinnedConnectionTaskExecutorTest, CancelRPC) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + + auto rcr = makeRCR(remote, BSONObj()); + auto pf = makePromiseFuture<void>(); + + // Schedule a command. + auto swCbHandle = pinnedTE->scheduleRemoteCommand( + std::move(rcr), [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + pf.promise.setWith([&] { return args.response.status; }); + }); + ASSERT_OK(swCbHandle.getStatus()); + auto cbHandle = swCbHandle.getValue(); + pinnedTE->cancel(cbHandle); + ASSERT_EQ(pf.future.getNoThrow(), TaskExecutor::kCallbackCanceledErrorStatus); + + pinnedTE->shutdown(); + pinnedTE->join(); +} + +TEST_F(PinnedConnectionTaskExecutorTest, ShutdownWithRPCInProgress) { + auto pinnedTE = makePinnedConnTaskExecutor(); + auto pf = makePromiseFuture<void>(); + ASSERT_OK(pinnedTE + ->scheduleRemoteCommand(makeRCR(HostAndPort("mock"), BSONObj()), + [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + pf.promise.setWith( + [&] { return args.response.status; }); + }) + .getStatus()); + pinnedTE->shutdown(); + ASSERT_EQ(pf.future.getNoThrow(), TaskExecutor::kCallbackCanceledErrorStatus); + pinnedTE->join(); +} + +TEST_F(PinnedConnectionTaskExecutorTest, CancelNonRPC) { + auto pinnedTE = makePinnedConnTaskExecutor(); + + auto pf = makePromiseFuture<void>(); + // Schedule some work + auto now = getNet()->now(); + auto swCbHandle = pinnedTE->scheduleWorkAt(now + Milliseconds(10), [&](auto&& cbArgs) { + pf.promise.setWith([&] { return cbArgs.status; }); + }); + + ASSERT_OK(swCbHandle.getStatus()); + auto cbHandle = swCbHandle.getValue(); + pinnedTE->cancel(cbHandle); + + ASSERT_EQ(pf.future.getNoThrow(), TaskExecutor::kCallbackCanceledErrorStatus); + + pinnedTE->shutdown(); + pinnedTE->join(); +} + +TEST_F(PinnedConnectionTaskExecutorTest, EnsureStreamIsUpdatedAfterUse) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + + auto rcr = makeRCR(remote, BSONObj()); + auto pf = makePromiseFuture<void>(); + // We haven't done any RPCs, so we shouldn't have touched any of the stream counters. + ASSERT_EQ(_indicateSuccessCalls.load(), 0); + ASSERT_EQ(_indicateUsedCalls.load(), 0); + ASSERT_EQ(_indicateFailureCalls.load(), 0); + + ASSERT_OK(pinnedTE + ->scheduleRemoteCommand(rcr, + [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + pf.promise.setWith( + [&] { return args.response.status; }); + }) + .getStatus()); + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyCameFromRequest(m, rcr); + return Status::OK(); + }); + expectSourceMessage([&]() { + auto message = makeOkReplyMessage(); + message.header().setResponseToMsgId(responseToId); + return message; + }); + + ASSERT_OK(pf.future.getNoThrow()); + + pinnedTE->shutdown(); + pinnedTE->join(); + + // We have compelted an RPC successfully using the leased stream: + ASSERT_EQ(_indicateSuccessCalls.load(), 1); + ASSERT_EQ(_indicateUsedCalls.load(), 1); + ASSERT_EQ(_indicateFailureCalls.load(), 0); +} + +TEST_F(PinnedConnectionTaskExecutorTest, StreamFailureShutsDownAndCancels) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + + // We haven't done any RPCs, so we shouldn't have touched any of the stream counters. + ASSERT_EQ(_indicateSuccessCalls.load(), 0); + ASSERT_EQ(_indicateUsedCalls.load(), 0); + ASSERT_EQ(_indicateFailureCalls.load(), 0); + + + // Schedule two RPCs + std::vector<Future<BSONObj>> results; + for (int i = 0; i < 2; ++i) { + auto promise = std::make_shared<Promise<BSONObj>>(NonNullPromiseTag{}); + results.push_back(promise->getFuture()); + auto extraFields = BSON("forTest" << i); + ASSERT_OK( + pinnedTE + ->scheduleRemoteCommand( + makeRCR(remote, extraFields), + [p = std::move(promise)](const TaskExecutor::RemoteCommandCallbackArgs& args) { + if (args.response.isOK()) { + p->emplaceValue(args.response.data); + } else { + p->setError(args.response.status); + } + }) + .getStatus()); + } + ASSERT_EQ(2, results.size()); + + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyAndDBName(m, BSON("hello" << 1), BSON("forTest" << 0), "admin"); + return Status::OK(); + }); + + // Fail the first request + Status testFailure{ErrorCodes::BadValue, "test failure"}; + expectSourceMessage([&]() { return testFailure; }); + auto localErr = results[0].getNoThrow().getStatus(); + ASSERT_EQ(localErr, testFailure); + + // The second should be cancelled automatically by shutdown. + ASSERT_EQ(results[1].getNoThrow(), TaskExecutor::kCallbackCanceledErrorStatus); + ASSERT(pinnedTE->isShuttingDown()); + + // We failed. + ASSERT_EQ(_indicateSuccessCalls.load(), 0); + ASSERT_EQ(_indicateUsedCalls.load(), 0); + ASSERT_EQ(_indicateFailureCalls.load(), 1); + pinnedTE->join(); +} + +/** + * We want to test the following sequence: + * (1) A command is scheduled. + * (2) The command fails due to a network error. + * (3) The command is notified of the failure (its onResponse callback is invoked). + * + * We want to ensure that the stream used by PCTE is destroyed _before_ the command is + * notified of the failure. This allows the underlying NetworkInterface to + * observe the failure on the initial stream & correctly update internally before it might + * be asked to provide another stream to i.e. retry the command. + */ +TEST_F(PinnedConnectionTaskExecutorTest, EnsureStreamDestroyedBeforeCommandCompleted) { + auto pinnedTE = makePinnedConnTaskExecutor(); + HostAndPort remote("mock"); + + auto rcr = makeRCR(remote, BSON("forTest" << 0)); + auto pf = makePromiseFuture<void>(); + ASSERT_EQ(_streamDestroyedCalls.load(), 0); + unittest::ThreadAssertionMonitor monitor; + auto completionCallback = [&](const TaskExecutor::RemoteCommandCallbackArgs& args) { + monitor.exec([&]() { + // Ensure the stream was destroyed before we are notified of the command completing. + ASSERT_EQ(_streamDestroyedCalls.load(), 1); + pf.promise.setWith([&] { return args.response.status; }); + monitor.notifyDone(); + }); + }; + + ASSERT_OK(pinnedTE->scheduleRemoteCommand(rcr, std::move(completionCallback)).getStatus()); + + int32_t responseToId; + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + assertMessageBodyAndDBName(m, BSON("hello" << 1), BSON("forTest" << 0), "admin"); + return Status::OK(); + }); + + // Fail the first request + Status testFailure{ErrorCodes::BadValue, "test failure"}; + expectSourceMessage([&]() { return testFailure; }); + + // Ensure we ran the completion callback. + monitor.wait(); + + auto localErr = pf.future.getNoThrow(); + ASSERT_EQ(localErr, testFailure); +} + +} // namespace +} // namespace mongo::executor diff --git a/src/mongo/executor/pinned_connection_task_executor_test_fixture.h b/src/mongo/executor/pinned_connection_task_executor_test_fixture.h new file mode 100644 index 00000000000..c5d2c77b524 --- /dev/null +++ b/src/mongo/executor/pinned_connection_task_executor_test_fixture.h @@ -0,0 +1,232 @@ +/** + * Copyright (C) 2023-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * 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 + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * 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 Server Side 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/base/status.h" +#include "mongo/db/dbmessage.h" +#include "mongo/executor/network_interface.h" +#include "mongo/executor/network_interface_mock.h" +#include "mongo/executor/pinned_connection_task_executor.h" +#include "mongo/executor/thread_pool_task_executor.h" +#include "mongo/executor/thread_pool_task_executor_test_fixture.h" +#include "mongo/rpc/op_msg_rpc_impls.h" +#include "mongo/transport/mock_session.h" +#include "mongo/transport/transport_layer.h" +#include "mongo/transport/transport_layer_mock.h" +#include "mongo/unittest/unittest.h" +#include "mongo/util/concurrency/thread_pool.h" + + +namespace mongo::executor { + + +class PinnedConnectionTaskExecutorTest : public ThreadPoolExecutorTest { + using SinkMessageCbT = std::function<Status(Message)>; + using SourceMessageCbT = std::function<StatusWith<Message>()>; + +public: + void setUp() override { + ThreadPoolExecutorTest::setUp(); + _session = std::make_shared<CustomMockSession>(this); + getNet()->setLeasedStreamMaker( + [this](HostAndPort hp) { return std::make_unique<LeasedStream>(hp, _session, this); }); + launchExecutorThread(); + } + + void tearDown() override { + ThreadPoolExecutorTest::tearDown(); + _session.reset(); + } + + Status sinkMessageCalled(Message message) { + stdx::unique_lock lk{_mutex}; + _hasWaitingSinkMessage = true; + _cv.wait(lk, [&] { return !!_sinkMessageExpectation || _isCanceled; }); + if (_isCanceled) { + // Consume the cancellation. + _isCanceled = false; + _sinkMessageExpectation = [&](auto&&) { return _cancellationError; }; + } + auto expectation = *std::exchange(_sinkMessageExpectation, {}); + _hasWaitingSinkMessage = false; + return expectation(message); + } + + void expectSinkMessage(SinkMessageCbT handler) { + stdx::lock_guard lk{_mutex}; + invariant(!_sinkMessageExpectation); + _sinkMessageExpectation = std::move(handler); + _cv.notify_one(); + } + + StatusWith<Message> sourceMessageCalled() { + stdx::unique_lock lk{_mutex}; + _hasWaitingSourceMessage = true; + _cv.wait(lk, [&] { return !!_sourceMessageExpectation || _isCanceled; }); + if (_isCanceled) { + // Consume the cancellation. + _isCanceled = false; + _sourceMessageExpectation = [&]() { return _cancellationError; }; + } + auto expectation = *std::exchange(_sourceMessageExpectation, {}); + _hasWaitingSourceMessage = false; + return expectation(); + } + + void expectSourceMessage(SourceMessageCbT handler) { + stdx::lock_guard lk{_mutex}; + invariant(!_sourceMessageExpectation); + _sourceMessageExpectation = std::move(handler); + _cv.notify_one(); + } + + void cancelAsyncOpsCalled() { + stdx::unique_lock lk{_mutex}; + _isCanceled = true; + _cv.notify_one(); + } + + bool hasReadyRequests() { + stdx::lock_guard lk{_mutex}; + return _hasWaitingSinkMessage || _hasWaitingSourceMessage; + } + + std::shared_ptr<PinnedConnectionTaskExecutor> makePinnedConnTaskExecutor() { + return std::make_shared<PinnedConnectionTaskExecutor>(getExecutorPtr(), getNet()); + } + +private: + std::shared_ptr<transport::Session> _session; + mutable Mutex _mutex; + stdx::condition_variable _cv; + boost::optional<SinkMessageCbT> _sinkMessageExpectation; + boost::optional<SourceMessageCbT> _sourceMessageExpectation; + bool _hasWaitingSinkMessage = false; + bool _hasWaitingSourceMessage = false; + bool _isCanceled = false; + Status _cancellationError = Status{ErrorCodes::SocketException, "Socket closed"}; + + class CustomMockSession : public transport::MockSessionBase { + public: + explicit CustomMockSession(PinnedConnectionTaskExecutorTest* fixture) : _fixture{fixture} {} + + transport::TransportLayer* getTransportLayer() const override { + return nullptr; + } + + void end() override { + *_connected = false; + } + + bool isConnected() override { + return *_connected; + } + + Status waitForData() noexcept override { + return Status::OK(); + } + + StatusWith<Message> sourceMessage() noexcept override { + return _fixture->sourceMessageCalled(); + } + + Status sinkMessage(Message message) noexcept override { + return _fixture->sinkMessageCalled(message); + } + + Future<void> asyncWaitForData() noexcept override { + return ExecutorFuture<void>(_fixture->getExecutorPtr()) + .then([this] { return waitForData(); }) + .unsafeToInlineFuture(); + } + + Future<Message> asyncSourceMessage(const BatonHandle& handle) noexcept override { + return ExecutorFuture<void>(_fixture->getExecutorPtr()) + .then([this] { return sourceMessage(); }) + .unsafeToInlineFuture(); + } + + Future<void> asyncSinkMessage(Message message, + const BatonHandle& handle) noexcept override { + return ExecutorFuture<void>(_fixture->getExecutorPtr()) + .then([this, m = std::move(message)] { return sinkMessage(m); }) + .unsafeToInlineFuture(); + } + + void cancelAsyncOperations(const BatonHandle& handle = nullptr) override { + _fixture->cancelAsyncOpsCalled(); + } + + private: + PinnedConnectionTaskExecutorTest* _fixture; + synchronized_value<bool> _connected{true}; + }; + + class LeasedStream : public NetworkInterface::LeasedStream { + public: + LeasedStream(HostAndPort hp, + std::shared_ptr<transport::Session> session, + PinnedConnectionTaskExecutorTest* fixture) + : _fixture{fixture} { + invariant(session); + _client = std::make_shared<AsyncDBClient>(hp, std::move(session), nullptr); + } + ~LeasedStream() { + _fixture->_streamDestroyedCalls.fetchAndAdd(1); + } + AsyncDBClient* getClient() override { + return _client.get(); + } + void indicateSuccess() override { + _fixture->_indicateSuccessCalls.fetchAndAdd(1); + } + void indicateUsed() override { + _fixture->_indicateUsedCalls.fetchAndAdd(1); + } + void indicateFailure(Status) override { + _fixture->_indicateFailureCalls.fetchAndAdd(1); + } + + private: + PinnedConnectionTaskExecutorTest* _fixture; + std::shared_ptr<AsyncDBClient> _client; + }; + +protected: + // Track the success/used/failure/destruction calls across LeasedStreams created via this + // fixture. Accessible to children so tests can read them directly. + AtomicWord<size_t> _indicateSuccessCalls{0}; + AtomicWord<size_t> _indicateUsedCalls{0}; + AtomicWord<size_t> _indicateFailureCalls{0}; + AtomicWord<size_t> _streamDestroyedCalls{0}; +}; + + +} // namespace mongo::executor diff --git a/src/mongo/executor/task_executor_cursor.cpp b/src/mongo/executor/task_executor_cursor.cpp index 46d1bd846b4..57a947ed4c6 100644 --- a/src/mongo/executor/task_executor_cursor.cpp +++ b/src/mongo/executor/task_executor_cursor.cpp @@ -36,29 +36,45 @@ #include "mongo/bson/bsonobjbuilder.h" #include "mongo/db/query/getmore_command_gen.h" #include "mongo/db/query/kill_cursors_gen.h" -#include "mongo/util/scopeguard.h" +#include "mongo/executor/pinned_connection_task_executor_factory.h" +#include "mongo/logv2/log.h" +#include "mongo/util/assert_util.h" #include "mongo/util/time_support.h" namespace mongo { namespace executor { +namespace { +MONGO_FAIL_POINT_DEFINE(blockBeforePinnedExecutorIsDestroyedOnUnderlying); +} // namespace -TaskExecutorCursor::TaskExecutorCursor(executor::TaskExecutor* executor, +TaskExecutorCursor::TaskExecutorCursor(std::shared_ptr<executor::TaskExecutor> executor, const RemoteCommandRequest& rcr, - Options&& options) - : _executor(executor), _rcr(rcr), _options(std::move(options)), _batchIter(_batch.end()) { + Options options) + : _rcr(rcr), _options(std::move(options)), _batchIter(_batch.end()) { if (rcr.opCtx) { _lsid = rcr.opCtx->getLogicalSessionId(); } + if (_options.pinConnection) { + _executor = makePinnedConnectionTaskExecutor(executor); + _underlyingExecutor = std::move(executor); + } else { + _executor = std::move(executor); + } _runRemoteCommand(_createRequest(_rcr.opCtx, _rcr.cmdObj)); } -TaskExecutorCursor::TaskExecutorCursor(executor::TaskExecutor* executor, +TaskExecutorCursor::TaskExecutorCursor(std::shared_ptr<executor::TaskExecutor> executor, + std::shared_ptr<executor::TaskExecutor> underlyingExec, CursorResponse&& response, RemoteCommandRequest& rcr, Options&& options) - : _executor(executor), _rcr(rcr), _options(std::move(options)), _batchIter(_batch.end()) { + : _executor(std::move(executor)), + _underlyingExecutor(std::move(underlyingExec)), + _rcr(rcr), + _options(std::move(options)), + _batchIter(_batch.end()) { tassert(6253101, "rcr must have an opCtx to use construct cursor from response", rcr.opCtx); _lsid = rcr.opCtx->getLogicalSessionId(); @@ -66,16 +82,16 @@ TaskExecutorCursor::TaskExecutorCursor(executor::TaskExecutor* executor, } TaskExecutorCursor::TaskExecutorCursor(TaskExecutorCursor&& other) - : _executor(other._executor), + : _executor(std::move(other._executor)), + _underlyingExecutor(std::move(other._underlyingExecutor)), _rcr(other._rcr), _options(std::move(other._options)), _lsid(other._lsid), - _cbHandle(std::move(other._cbHandle)), + _cmdState(std::move(other._cmdState)), _cursorId(other._cursorId), _millisecondsWaiting(other._millisecondsWaiting), _ns(other._ns), _batchNum(other._batchNum), - _pipe(std::move(other._pipe)), _additionalCursors(std::move(other._additionalCursors)) { // Copy the status of the batch. auto batchIterIndex = other._batchIter - other._batch.begin(); @@ -91,30 +107,74 @@ TaskExecutorCursor::TaskExecutorCursor(TaskExecutorCursor&& other) } // Other is no longer responsible for this cursor id. other._cursorId = 0; - // Other should not cancel the callback on destruction. - other._cbHandle = boost::none; + + // Other no longer owns the state for the in progress command (if there is any). + other._cmdState.reset(); } TaskExecutorCursor::~TaskExecutorCursor() { try { - if (_cbHandle) { - _executor->cancel(*_cbHandle); + if (_cursorId < kMinLegalCursorId || _options.pinConnection) { + // The initial find to establish the cursor has to be canceled to avoid leaking cursors. + // Once the cursor is established, killing the cursor will interrupt any ongoing + // `getMore` operation. + // Additionally, in pinned mode, we should cancel any in-progress RPC if there is one, + // even at the cost of churning the connection, because it's the only way to interrupt + // the ongoing operation. + if (_cmdState) { + _executor->cancel(_cmdState->cbHandle); + } + if (_cursorId < kMinLegalCursorId) { + return; + } } - if (_cursorId >= kMinLegalCursorId) { - // We deliberately ignore failures to kill the cursor. This "best effort" is acceptable - // because some timeout mechanism on the remote host can be expected to reap it later. - // - // That timeout mechanism could be the default cursor timeout, or the logical session - // timeout if an lsid is used. - _executor - ->scheduleRemoteCommand( - _createRequest(nullptr, - KillCursorsCommandRequest(_ns, {_cursorId}).toBSON(BSONObj{})), - [](const auto&) {}) - .isOK(); + // We deliberately ignore failures to kill the cursor. This "best effort" is acceptable + // because some timeout mechanism on the remote host can be expected to reap it later. + // + // That timeout mechanism could be the default cursor timeout, or the logical session + // timeout if an lsid is used. + // + // In non-pinned mode, killing the cursor also interrupts any ongoing getMore operations on + // this cursor. Avoid canceling the remote command through its callback handle as that may + // close the underlying connection. + // + // In pinned mode, we must await completion of the killCursors to safely reuse the pinned + // connection. This requires allocating an executor thread (from `_underlyingExecutor`) upon + // completion of the killCursors command to shutdown and destroy the pinned executor. This + // is necessary as joining an executor from its own threads results in a deadlock. + TaskExecutor::RemoteCommandCallbackFn callbackToRun = [](const auto&) {}; + if (_options.pinConnection) { + invariant(_underlyingExecutor, + "TaskExecutorCursor in pinning mode must have an underlying executor"); + callbackToRun = [main = _executor, underlying = _underlyingExecutor](const auto&) { + underlying->schedule([main = std::move(main)](const auto&) { + if (MONGO_unlikely( + blockBeforePinnedExecutorIsDestroyedOnUnderlying.shouldFail())) { + LOGV2(7361300, + "Hanging before destroying a TaskExecutorCursor's pinning executor."); + blockBeforePinnedExecutorIsDestroyedOnUnderlying.pauseWhileSet(); + } + // Returning from this callback will destroy the pinned executor on + // underlying if this is the last TaskExecutorCursor using that pinned executor. + }); + }; + } + auto swCallback = _executor->scheduleRemoteCommand( + _createRequest(nullptr, KillCursorsCommandRequest(_ns, {_cursorId}).toBSON(BSONObj{})), + callbackToRun); + + // It's possible the executor is already shutdown and rejects work. If so, run the callback + // inline. + if (!swCallback.isOK()) { + TaskExecutor::RemoteCommandCallbackArgs args( + _executor.get(), {}, {}, swCallback.getStatus()); + callbackToRun(args); } - } catch (const DBException&) { + } catch (const DBException& ex) { + LOGV2(6531704, + "Encountered an error while destroying a cursor executor", + "error"_attr = ex.toStatus()); } } @@ -136,7 +196,7 @@ void TaskExecutorCursor::populateCursor(OperationContext* opCtx) { _cursorId == kUnitializedCursorId); tassert(6253503, "populateCursors should only be called after a remote command has been run", - _cbHandle); + _cmdState); // We really only care about populating the cursor "first batch" fields, but at some point we'll // have to do all of the work done by this function anyway. This would have been called by // getNext() the first time it was called. @@ -167,21 +227,18 @@ const RemoteCommandRequest& TaskExecutorCursor::_createRequest(OperationContext* } void TaskExecutorCursor::_runRemoteCommand(const RemoteCommandRequest& rcr) { - _cbHandle = uassertStatusOK(_executor->scheduleRemoteCommand( - rcr, [p = _pipe.producer](const TaskExecutor::RemoteCommandCallbackArgs& args) { - try { - if (args.response.isOK()) { - p.push(args.response.data); - } else { - p.push(args.response.status); - } - } catch (const DBException&) { - // If anything goes wrong, make sure we close the pipe to wake the caller of - // getNext() - p.close(); + auto state = std::make_shared<CommandState>(); + state->cbHandle = uassertStatusOK(_executor->scheduleRemoteCommand( + rcr, [state](const TaskExecutor::RemoteCommandCallbackArgs& args) { + if (args.response.isOK()) { + state->promise.emplaceValue(args.response.data); + } else { + state->promise.setError(args.response.status); } })); + _cmdState.swap(state); } + void TaskExecutorCursor::_processResponse(OperationContext* opCtx, CursorResponse&& response) { // If this was our first batch. if (_cursorId == kUnitializedCursorId) { @@ -196,23 +253,51 @@ void TaskExecutorCursor::_processResponse(OperationContext* opCtx, CursorRespons _batch = response.releaseBatch(); _batchIter = _batch.begin(); - // If we got a cursor id back, pre-fetch the next batch - if (_cursorId) { - GetMoreCommandRequest getMoreRequest(_cursorId, _ns.coll().toString()); - getMoreRequest.setBatchSize(_options.batchSize); + // If the previous response contained a cursorId and pre-fetching is enabled, schedule the + // getMore. + if ((_cursorId != kClosedCursorId) && _options.preFetchNextBatch) { + _scheduleGetMore(opCtx); + } +} + +void TaskExecutorCursor::_scheduleGetMore(OperationContext* opCtx) { + // The previous response must have returned an open cursor ID. + invariant(_cursorId >= kMinLegalCursorId); + // There cannot be an existing in-flight request. + invariant(!_cmdState); + GetMoreCommandRequest getMoreRequest(_cursorId, _ns.coll().toString()); + getMoreRequest.setBatchSize(_options.batchSize); + + if (_options.getMoreAugmentationWriter) { + // Prefetching must be disabled to use the augmenting functionality. + invariant(!_options.preFetchNextBatch); + BSONObjBuilder getMoreBob; + getMoreRequest.serialize({}, &getMoreBob); + _options.getMoreAugmentationWriter(getMoreBob); + _runRemoteCommand(_createRequest(opCtx, getMoreBob.obj())); + } else { _runRemoteCommand(_createRequest(opCtx, getMoreRequest.toBSON({}))); } } void TaskExecutorCursor::_getNextBatch(OperationContext* opCtx) { - invariant(_cbHandle, "_getNextBatch() requires an async request to have already been sent."); + // If we don't have an in-flight request, schedule one. This will occur when the + // 'preFetchNextBatch' option is false. + if (!_cmdState) { + invariant(!_options.preFetchNextBatch); + _scheduleGetMore(opCtx); + } + + // There should be an in-flight request at this point, either sent asyncronously when we + // processed the previous response or just scheduled. + invariant(_cmdState); invariant(_cursorId != kClosedCursorId); auto clock = opCtx->getServiceContext()->getPreciseClockSource(); auto dateStart = clock->now(); // pull out of the pipe before setting cursor id so we don't spoil this object if we're opCtx // interrupted - auto out = _pipe.consumer.pop(opCtx); + auto out = _cmdState->promise.getFuture().getNoThrow(opCtx); auto dateEnd = clock->now(); _millisecondsWaiting += std::max(Milliseconds(0), dateEnd - dateStart); uassertStatusOK(out); @@ -226,20 +311,31 @@ void TaskExecutorCursor::_getNextBatch(OperationContext* opCtx) { // if we've received a response from our last request (initial or getmore), our remote operation // is done. - _cbHandle.reset(); + _cmdState.reset(); // Parse into a vector in case the remote sent back multiple cursors. auto cursorResponses = CursorResponse::parseFromBSONMany(out.getValue()); tassert(6253100, "Expected at least one response for cursor", cursorResponses.size() > 0); CursorResponse cr = uassertStatusOK(std::move(cursorResponses[0])); _processResponse(opCtx, std::move(cr)); + // If we have more responses, build them into cursors then hold them until a caller accesses // them. Skip the first response, we used it to populate this cursor. + // Ensure we update the RCR we give to each 'child cursor' with the current opCtx. + auto freshRcr = _createRequest(opCtx, _rcr.cmdObj); + auto copyOptions = [&] { + TaskExecutorCursor::Options options; + // In the case that pinConnection is true, we need to ensure that additional cursors also + // pin their connection to the same socket as the original cursor. + options.pinConnection = _options.pinConnection; + return options; + }; for (unsigned int i = 1; i < cursorResponses.size(); ++i) { _additionalCursors.emplace_back(_executor, + _underlyingExecutor, uassertStatusOK(std::move(cursorResponses[i])), - _rcr, - TaskExecutorCursor::Options()); + freshRcr, + copyOptions()); } } diff --git a/src/mongo/executor/task_executor_cursor.h b/src/mongo/executor/task_executor_cursor.h index b18ea481c88..6de3a7664c1 100644 --- a/src/mongo/executor/task_executor_cursor.h +++ b/src/mongo/executor/task_executor_cursor.h @@ -30,6 +30,7 @@ #pragma once #include <boost/optional.hpp> +#include <memory> #include <vector> #include "mongo/base/status_with.h" @@ -40,9 +41,9 @@ #include "mongo/db/query/cursor_response.h" #include "mongo/executor/remote_command_request.h" #include "mongo/executor/task_executor.h" +#include "mongo/executor/task_executor_cursor_parameters_gen.h" #include "mongo/util/duration.h" -#include "mongo/util/net/hostandport.h" -#include "mongo/util/producer_consumer_queue.h" +#include "mongo/util/future.h" namespace mongo { namespace executor { @@ -53,8 +54,8 @@ namespace executor { * * The main differentiator for this type over DBClientCursor is the use of a task executor (which * provides access to a different connection pool, as well as interruptibility) and the ability to - * overlap getMores. This starts fetching the next batch as soon as one is exhausted (rather than - * on a call to getNext()). + * overlap getMores. This starts fetching the next batch as soon as the previous one is received + * (rather than on a call to 'getNext()'). */ class TaskExecutorCursor { public: @@ -68,27 +69,44 @@ public: struct Options { boost::optional<int64_t> batchSize; + bool pinConnection{gPinTaskExecCursorConns.load()}; + // If true, we will fetch the next batch as soon as the current one is recieved. + // If false, we will fetch the next batch when the current batch is exhausted and + // 'getNext()' is invoked. + bool preFetchNextBatch{true}; + + // This function, if specified, may modify a getMore request to include additional + // information. + std::function<void(BSONObjBuilder& bob)> getMoreAugmentationWriter; + + Options() {} }; /** - * Construct the cursor with a RemoteCommandRequest wrapping the initial command + * Construct the cursor with a RemoteCommandRequest wrapping the initial command. + * + * Doesn't retry the command if we fail to establish the cursor. To create a TaskExecutorCursor + * with the option to retry the initial command, see `makeTaskExecutorCursor`below. * * One value is carried over in successive calls to getMore/killCursor: * * opCtx - The Logical Session Id from the initial command is carried over in all later stages. * NOTE - the actual command must not include the lsid */ - explicit TaskExecutorCursor(executor::TaskExecutor* executor, - const RemoteCommandRequest& rcr, - Options&& options = {}); + TaskExecutorCursor(std::shared_ptr<executor::TaskExecutor> executor, + const RemoteCommandRequest& rcr, + Options options = {}); /** * Construct the cursor from a cursor response from a previously executed RemoteCommandRequest. * The executor is used for subsequent getMore calls. Uses the original RemoteCommandRequest * to build subsequent commands. Takes ownership of the CursorResponse and gives it to the new * cursor. + * If the cursor should reuse the original transport connection that opened the original + * cursor, make sure the pinning executor that was used to open that cursor is provided. */ - TaskExecutorCursor(executor::TaskExecutor* executor, + TaskExecutorCursor(std::shared_ptr<executor::TaskExecutor> executor, + std::shared_ptr<executor::TaskExecutor> underlyingExec, CursorResponse&& response, RemoteCommandRequest& rcr, Options&& options = {}); @@ -158,14 +176,6 @@ public: return _additionalCursors.size(); } - /** - * Return the callback that this cursor is waiting on. Can be used to block on getting a - * response to this request. Can be boost::none. - */ - auto getCallbackHandle() { - return _cbHandle; - } - private: /** * Runs a remote command and pipes the output back to this object @@ -173,24 +183,34 @@ private: void _runRemoteCommand(const RemoteCommandRequest& rcr); /** - * Gets the next batch with interruptibility via the opCtx + * Gets the next batch with interruptibility via the opCtx. */ void _getNextBatch(OperationContext* opCtx); /** * Helper for '_getNextBatch' that handles the reading of the 'CursorResponse' object and - * storing of relevant values. This is also responsible for issuing a getMore request if it - * is required to populate the next batch. + * storing of relevant values. */ void _processResponse(OperationContext* opCtx, CursorResponse&& response); - /** * Create a new request, annotating with lsid and current opCtx */ const RemoteCommandRequest& _createRequest(OperationContext* opCtx, const BSONObj& cmd); - executor::TaskExecutor* _executor; + /** + * Schedules a 'GetMore' request to run asyncronously. + * This function can only be invoked when: + * - There is no in-flight request ('_cmdState' is null). + * - We have an open '_cursorId'. + */ + void _scheduleGetMore(OperationContext* opCtx); + + std::shared_ptr<executor::TaskExecutor> _executor; + // If we are pinning connections, we need to keep a separate reference to the + // non-pinning, normal executor, so that we can shut down the pinned executor + // out-of-line. + std::shared_ptr<executor::TaskExecutor> _underlyingExecutor; // Used as a scratch pad for the successive scheduleRemoteCommand calls RemoteCommandRequest _rcr; @@ -200,8 +220,20 @@ private: // If the opCtx is in our initial request, re-use it for all subsequent operations boost::optional<LogicalSessionId> _lsid; - // Stash the callbackhandle for the current outstanding operation - boost::optional<TaskExecutor::CallbackHandle> _cbHandle; + struct CommandState { + TaskExecutor::CallbackHandle cbHandle; + SharedPromise<BSONObj> promise; + }; + + /** + * Maintains the state for the in progress command (if there is any): + * - Handle for the task scheduled on `_executor`. + * - A promise that will be emplaced by the result of running the command. + * + * The state may outlive `TaskExecutorCursor` and is shared with the callback that runs on + * `_executor` upon completion of the remote command. + */ + std::shared_ptr<CommandState> _cmdState; CursorId _cursorId = kUnitializedCursorId; @@ -222,13 +254,33 @@ private: decltype(_batch)::iterator _batchIter; long long _batchNum = 0; - // Multi producer because we hold onto the producer side in this object, as well as placing it - // into callbacks for the task executor - MultiProducerSingleConsumerQueue<StatusWith<BSONObj>>::Pipe _pipe; - // Cursors built from the responses returned alongside the results for this cursor. std::vector<TaskExecutorCursor> _additionalCursors; }; +// Make a new TaskExecutorCursor using the provided executor, RCR, and options. If we fail to create +// the cursor, the retryPolicy can inspect the error and make a decision as to whether we should +// retry. If we do retry, the error is swallowed and another attempt is made. If we don't retry, +// this function throws the error we failed with. +inline TaskExecutorCursor makeTaskExecutorCursor( + OperationContext* opCtx, + std::shared_ptr<executor::TaskExecutor> executor, + const RemoteCommandRequest& rcr, + TaskExecutorCursor::Options options = {}, + std::function<bool(Status)> retryPolicy = nullptr) { + for (;;) { + try { + TaskExecutorCursor tec(executor, rcr, options); + tec.populateCursor(opCtx); + return tec; + } catch (const DBException& ex) { + bool shouldRetry = retryPolicy && retryPolicy(ex.toStatus()); + if (!shouldRetry) { + throw; + } + } + } +} + } // namespace executor } // namespace mongo diff --git a/src/mongo/executor/task_executor_cursor_integration_test.cpp b/src/mongo/executor/task_executor_cursor_integration_test.cpp index e3c3fde671d..8a8433afd72 100644 --- a/src/mongo/executor/task_executor_cursor_integration_test.cpp +++ b/src/mongo/executor/task_executor_cursor_integration_test.cpp @@ -27,6 +27,8 @@ * it in the license file. */ +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest + #include "mongo/platform/basic.h" #include "mongo/executor/task_executor_cursor.h" @@ -34,9 +36,12 @@ #include "mongo/client/dbclient_base.h" #include "mongo/db/concurrency/locker_noop_client_observer.h" #include "mongo/db/namespace_string.h" +#include "mongo/executor/connection_pool_stats.h" #include "mongo/executor/network_interface_factory.h" #include "mongo/executor/network_interface_thread_pool.h" +#include "mongo/executor/pinned_connection_task_executor.h" #include "mongo/executor/thread_pool_task_executor.h" +#include "mongo/logv2/log.h" #include "mongo/unittest/integration_test.h" #include "mongo/unittest/unittest.h" @@ -51,48 +56,75 @@ public: } void setUp() override { - std::shared_ptr<NetworkInterface> ni = makeNetworkInterface("TaskExecutorCursorTest"); - auto tp = std::make_unique<NetworkInterfaceThreadPool>(ni.get()); + _ni = makeNetworkInterface("TaskExecutorCursorTest"); + auto tp = std::make_unique<NetworkInterfaceThreadPool>(_ni.get()); - _executor = std::make_unique<ThreadPoolTaskExecutor>(std::move(tp), std::move(ni)); + _executor = std::make_shared<ThreadPoolTaskExecutor>(std::move(tp), _ni); _executor->startup(); }; void tearDown() override { _executor->shutdown(); + _executor->join(); _executor.reset(); }; - TaskExecutor* executor() { - return _executor.get(); + std::shared_ptr<TaskExecutor> executor() { + return _executor; } - ServiceContext::UniqueServiceContext _serviceCtx = ServiceContext::make(); - std::unique_ptr<ThreadPoolTaskExecutor> _executor; -}; + auto net() { + return _ni.get(); + } + + auto makeOpCtx() { + return _client->makeOperationContext(); + } + TaskExecutor::CallbackHandle scheduleRemoteCommand(OperationContext* opCtx, + HostAndPort target, + BSONObj cmd) { + LOGV2(6531702, "About to run a remote command", "cmd"_attr = cmd); + RemoteCommandRequest rcr(target, "admin", cmd, opCtx); + auto swHandle = executor()->scheduleRemoteCommand( + std::move(rcr), [](const TaskExecutor::RemoteCommandCallbackArgs&) {}); + return uassertStatusOK(swHandle); + } -// Test that we can actually use a TaskExecutorCursor to read multiple batches from a remote host -TEST_F(TaskExecutorCursorFixture, Basic) { - auto client = _serviceCtx->makeClient("TaskExecutorCursorTest"); - auto opCtx = client->makeOperationContext(); + void runRemoteCommand(OperationContext* opCtx, HostAndPort target, BSONObj cmd) { + auto cbHandle = scheduleRemoteCommand(opCtx, std::move(target), cmd); + executor()->wait(cbHandle, opCtx); + LOGV2(6531703, "Finished running remote command", "cmd"_attr = cmd); + } + +private: + ServiceContext::UniqueServiceContext _serviceCtx = ServiceContext::make(); + std::shared_ptr<ThreadPoolTaskExecutor> _executor; + std::shared_ptr<NetworkInterface> _ni; + ServiceContext::UniqueClient _client = _serviceCtx->makeClient("TaskExecutorCursorTest"); +}; - // Write 100 documents to "test.test" via dbclient +size_t createTestData(std::string ns, size_t numDocs) { auto swConn = unittest::getFixtureConnectionString().connect("TaskExecutorCursorTest"); uassertStatusOK(swConn.getStatus()); auto dbclient = std::move(swConn.getValue()); - const size_t numDocs = 100; - std::vector<BSONObj> docs; docs.reserve(numDocs); for (size_t i = 0; i < numDocs; ++i) { docs.emplace_back(BSON("x" << int(i))); } - dbclient->dropCollection("test.test"); - dbclient->insert("test.test", docs); - ASSERT_EQUALS(dbclient->count(NamespaceString("test.test")), numDocs); + dbclient->dropCollection(ns); + dbclient->insert(ns, docs); + return dbclient->count(NamespaceString(ns)); +} +// Test that we can actually use a TaskExecutorCursor to read multiple batches from a remote host. +TEST_F(TaskExecutorCursorFixture, Basic) { + const size_t numDocs = 100; + ASSERT_EQ(createTestData("test.test", numDocs), numDocs); + + auto opCtx = makeOpCtx(); RemoteCommandRequest rcr(unittest::getFixtureConnectionString().getServers().front(), "test", BSON("find" @@ -114,6 +146,177 @@ TEST_F(TaskExecutorCursorFixture, Basic) { ASSERT_EQUALS(count, numDocs); } +// Test that we can actually use a TaskExecutorCursor that pins it's connection to read multiple +// batches from a remote host. +TEST_F(TaskExecutorCursorFixture, BasicPinned) { + const size_t numDocs = 100; + ASSERT_EQ(createTestData("test.test", numDocs), numDocs); + + auto opCtx = makeOpCtx(); + RemoteCommandRequest rcr(unittest::getFixtureConnectionString().getServers().front(), + "test", + BSON("find" + << "test" + << "batchSize" << 10), + opCtx.get()); + + TaskExecutorCursor tec(executor(), rcr, [this] { + TaskExecutorCursor::Options opts; + opts.batchSize = 10; + opts.pinConnection = true; + return opts; + }()); + + size_t count = 0; + while (auto doc = tec.getNext(opCtx.get())) { + count++; + } + + ASSERT_EQUALS(count, numDocs); +} + +// Test that when a TaskExecutorCursor is used in pinning-mode, the pinned executor's destruction +// is scheduled on the underlying executor. +TEST_F(TaskExecutorCursorFixture, PinnedExecutorDestroyedOnUnderlying) { + const size_t numDocs = 100; + ASSERT_EQ(createTestData("test.test", numDocs), numDocs); + + auto opCtx = makeOpCtx(); + RemoteCommandRequest rcr(unittest::getFixtureConnectionString().getServers().front(), + "test", + BSON("find" + << "test" + << "batchSize" << 10), + opCtx.get()); + + boost::optional<TaskExecutorCursor> tec; + tec.emplace(executor(), rcr, [] { + TaskExecutorCursor::Options opts; + opts.batchSize = 10; + opts.pinConnection = true; + return opts; + }()); + // Fetch a documents to make sure the TEC was initialized properly. + ASSERT(tec->getNext(opCtx.get())); + // Enable the failpoint in the integration test process. + { + FailPointEnableBlock fpb("blockBeforePinnedExecutorIsDestroyedOnUnderlying"); + auto initialTimesEntered = fpb.initialTimesEntered(); + // Destroy the TEC and ensure we reach the code block that will destroy the pinned executor. + tec.reset(); + LOGV2(7361301, "Waiting for TaskExecutorCursor to destroy its pinning executor."); + fpb->waitForTimesEntered(initialTimesEntered + 1); + } + // Allow the pinned executor's destruction to proceed. +} + +/** + * Verifies that the underlying connection used to run `getMore` commands remains open, even after + * the instance of `TaskExecutorCursor` is destroyed. + * + * The test goes through the following steps: + * - Load test data into "test.test". + * - Make sure there are enough connections to run the test. + * - Configure a fail-point to block the initial `getMore` that populates the cursor. + * - Create an instance of `TaskExecutorCursor`. The executor will send out an asynchronous + * `getMore` to populate the cursor, but the command does not return so long as the fail-point + * is enabled. + * - Count the total number of connections available before destroying `TaskExecutorCursor`. + * - Destroy the instance of `TaskExecutorCursor`, disable the fail-point, and wait for all + * connections to become idle. + * - Recount the number of connections and verify that no connection is closed. + * + * See SERVER-65317 for more context. + */ +TEST_F(TaskExecutorCursorFixture, ConnectionRemainsOpenAfterKillingTheCursor) { + const size_t numDocs = 100; + ASSERT_EQ(createTestData("test.test", numDocs), numDocs); + + auto opCtx = makeOpCtx(); + const auto target = unittest::getFixtureConnectionString().getServers().front(); + + auto getConnectionStatsForTarget = [&] { + ConnectionPoolStats stats; + executor()->appendConnectionStats(&stats); + return stats.statsByHost[target]; + }; + + // We only need at most four connections to run this test, which will run the following + // commands: `find`, `getMore`, `killCursor`, and `configureFailPoint`. Thus, the rest of this + // test won't need to create new connections and the number of open connections should remain + // unchanged. + const size_t kNumConnections = 4; + std::vector<TaskExecutor::CallbackHandle> handles; + auto cmd = BSON("find" + << "test" + << "filter" + << BSON("$where" + << "sleep(100); return true;")); + for (size_t i = 0; i < kNumConnections; i++) { + handles.emplace_back(scheduleRemoteCommand(opCtx.get(), target, cmd)); + } + for (auto cbHandle : handles) { + executor()->wait(cbHandle); + } + + ConnectionStatsPer beforeStats; + { + const auto fpName = "waitAfterCommandFinishesExecution"; + runRemoteCommand(opCtx.get(), + target, + BSON("configureFailPoint" + << fpName << "mode" + << "alwaysOn" + << "data" + << BSON("ns" + << "test.test" + << "commands" << BSON_ARRAY("getMore")))); + ScopeGuard guard([&] { + runRemoteCommand(opCtx.get(), + target, + BSON("configureFailPoint" << fpName << "mode" + << "off")); + }); + + RemoteCommandRequest rcr(target, + "test", + BSON("find" + << "test" + << "batchSize" << 10), + opCtx.get()); + + TaskExecutorCursor tec(executor(), rcr, [] { + TaskExecutorCursor::Options opts; + opts.batchSize = 10; + return opts; + }()); + + tec.populateCursor(opCtx.get()); + + // At least one of the connections is busy running the initial `getMore` command to populate + // the cursor. The command is blocked on the remote host and does not return until after the + // destructor for `tec` returns. + beforeStats = getConnectionStatsForTarget(); + ASSERT_GTE(beforeStats.inUse, 1); + } + + // Wait for all connections to become idle -- this ensures all tasks scheduled as part of + // cleaning up `tec` have finished running. + while (getConnectionStatsForTarget().inUse > 0) { + LOGV2(6531701, "Waiting for all connections to become idle"); + sleepFor(Seconds(1)); + } + + const auto afterStats = getConnectionStatsForTarget(); + auto countOpenConns = [](const ConnectionStatsPer& stats) { + return stats.inUse + stats.available + stats.refreshing + stats.leased; + }; + + // Verify that no connection is created or closed. + ASSERT_EQ(beforeStats.created, afterStats.created); + ASSERT_EQ(countOpenConns(beforeStats), countOpenConns(afterStats)); +} + } // namespace } // namespace executor } // namespace mongo diff --git a/src/mongo/executor/task_executor_cursor_parameters.idl b/src/mongo/executor/task_executor_cursor_parameters.idl new file mode 100644 index 00000000000..3ae7b73429e --- /dev/null +++ b/src/mongo/executor/task_executor_cursor_parameters.idl @@ -0,0 +1,41 @@ +# Copyright (C) 2023-present MongoDB, Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the Server Side Public License, version 1, +# as published by MongoDB, Inc. +# +# 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 +# Server Side Public License for more details. +# +# You should have received a copy of the Server Side Public License +# along with this program. If not, see +# <http://www.mongodb.com/licensing/server-side-public-license>. +# +# 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 Server Side 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. +# + +global: + cpp_namespace: "mongo" + +server_parameters: + pinTaskExecCursorConns: + description: >- + TaskExecutorCursor will `pin` the transport connection it uses to establish + cursors and re-use that same connection for all subsequent operations on those cursors. + set_at: [ startup, runtime ] + cpp_vartype: AtomicWord<bool> + cpp_varname: "gPinTaskExecCursorConns" + default: false + diff --git a/src/mongo/executor/task_executor_cursor_test.cpp b/src/mongo/executor/task_executor_cursor_test.cpp index 09d1a1c301b..b99edbf715d 100644 --- a/src/mongo/executor/task_executor_cursor_test.cpp +++ b/src/mongo/executor/task_executor_cursor_test.cpp @@ -32,65 +32,564 @@ #include "mongo/platform/basic.h" #include "mongo/db/concurrency/locker_noop_client_observer.h" +#include "mongo/executor/pinned_connection_task_executor.h" +#include "mongo/executor/pinned_connection_task_executor_test_fixture.h" #include "mongo/executor/task_executor_cursor.h" #include "mongo/executor/thread_pool_task_executor_test_fixture.h" +#include "mongo/rpc/op_msg_rpc_impls.h" #include "mongo/unittest/bson_test_util.h" +#include "mongo/unittest/thread_assertion_monitor.h" #include "mongo/unittest/unittest.h" namespace mongo { namespace executor { namespace { +BSONObj buildCursorResponse(StringData fieldName, size_t start, size_t end, size_t cursorId) { + BSONObjBuilder bob; + { + BSONObjBuilder cursor(bob.subobjStart("cursor")); + { + BSONArrayBuilder batch(cursor.subarrayStart(fieldName)); + + for (size_t i = start; i <= end; ++i) { + BSONObjBuilder doc(batch.subobjStart()); + doc.append("x", int(i)); + } + } + cursor.append("id", (long long)(cursorId)); + cursor.append("ns", "test.test"); + } + bob.append("ok", int(1)); + return bob.obj(); +} + +BSONObj buildMultiCursorResponse(StringData fieldName, + size_t start, + size_t end, + std::vector<size_t> cursorIds) { + BSONObjBuilder bob; + { + BSONArrayBuilder cursors; + int baseCursorValue = 1; + for (auto cursorId : cursorIds) { + BSONObjBuilder cursor; + BSONArrayBuilder batch; + ASSERT(start < end && end < INT_MAX); + for (size_t i = start; i <= end; ++i) { + batch.append(BSON("x" << static_cast<int>(i) * baseCursorValue).getOwned()); + } + cursor.append(fieldName, batch.arr()); + cursor.append("id", (long long)(cursorId)); + cursor.append("ns", "test.test"); + auto cursorObj = BSON("cursor" << cursor.done() << "ok" << 1); + cursors.append(cursorObj.getOwned()); + ++baseCursorValue; + } + bob.append("cursors", cursors.arr()); + } + bob.append("ok", 1); + return bob.obj(); +} + /** * Fixture for the task executor cursor tests which offers some convenience methods to help with - * scheduling responses + * scheduling responses. Uses the CRTP pattern so that the tests can be shared between child-classes + * that provide their own implementations of the network-mocking needed for the tests. */ -class TaskExecutorCursorFixture : public ThreadPoolExecutorTest { +template <typename Derived, typename Base> +class TaskExecutorCursorTestFixture : public Base { public: - TaskExecutorCursorFixture() { + TaskExecutorCursorTestFixture() { serviceCtx->registerClientObserver(std::make_unique<LockerNoopClientObserver>()); } void setUp() override { - ThreadPoolExecutorTest::setUp(); - + Base::setUp(); client = serviceCtx->makeClient("TaskExecutorCursorTest"); opCtx = client->makeOperationContext(); - - launchExecutorThread(); + static_cast<Derived*>(this)->postSetUp(); } void tearDown() override { opCtx.reset(); client.reset(); - ThreadPoolExecutorTest::tearDown(); + Base::tearDown(); } BSONObj scheduleSuccessfulCursorResponse(StringData fieldName, size_t start, size_t end, size_t cursorId) { - NetworkInterfaceMock::InNetworkGuard ing(getNet()); + return static_cast<Derived*>(this)->scheduleSuccessfulCursorResponse( + fieldName, start, end, cursorId); + } - BSONObjBuilder bob; - { - BSONObjBuilder cursor(bob.subobjStart("cursor")); - { - BSONArrayBuilder batch(cursor.subarrayStart(fieldName)); - - for (size_t i = start; i <= end; ++i) { - BSONObjBuilder doc(batch.subobjStart()); - doc.append("x", int(i)); - } + BSONObj scheduleSuccessfulMultiCursorResponse(StringData fieldName, + size_t start, + size_t end, + std::vector<size_t> cursorIds) { + return static_cast<Derived*>(this)->scheduleSuccessfulMultiCursorResponse( + fieldName, start, end, cursorIds); + } + + void scheduleErrorResponse(Status error) { + return static_cast<Derived*>(this)->scheduleErrorResponse(error); + } + void blackHoleNextOutgoingRequest() { + return static_cast<Derived*>(this)->blackHoleNextOutgoingRequest(); + } + + BSONObj scheduleSuccessfulKillCursorResponse(size_t cursorId) { + return static_cast<Derived*>(this)->scheduleSuccessfulKillCursorResponse(cursorId); + } + + TaskExecutorCursor makeTec(RemoteCommandRequest rcr, + TaskExecutorCursor::Options&& options = {}) { + return static_cast<Derived*>(this)->makeTec(rcr, std::move(options)); + } + + bool hasReadyRequests() { + return static_cast<Derived*>(this)->hasReadyRequests(); + } + + Base& asBase() { + return *this; + } + + /** + * Ensure we work for a single simple batch + */ + void SingleBatchWorksTest() { + const auto findCmd = BSON("find" + << "test" + << "batchSize" << 2); + const CursorId cursorId = 0; + + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr); + + ASSERT_BSONOBJ_EQ(findCmd, scheduleSuccessfulCursorResponse("firstBatch", 1, 2, cursorId)); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + ASSERT_FALSE(hasReadyRequests()); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + ASSERT_FALSE(tec.getNext(opCtx.get())); + } + + /** + * Ensure the firstBatch can be read correctly when multiple cursors are returned. + */ + void MultipleCursorsSingleBatchSucceedsTest() { + const auto aggCmd = BSON("aggregate" + << "test" + << "pipeline" + << BSON_ARRAY(BSON("returnMultipleCursors" << true))); + + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", aggCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr); + + ASSERT_BSONOBJ_EQ(aggCmd, + scheduleSuccessfulMultiCursorResponse("firstBatch", 1, 2, {0, 0})); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + ASSERT_FALSE(tec.getNext(opCtx.get())); + + auto cursorVec = tec.releaseAdditionalCursors(); + ASSERT_EQUALS(cursorVec.size(), 1); + auto secondCursor = std::move(cursorVec[0]); + + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 2); + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 4); + ASSERT_FALSE(hasReadyRequests()); + + ASSERT_FALSE(secondCursor.getNext(opCtx.get())); + } + /** + * The operation context under which we send the original cursor-establishing command + * can be destructed before getNext is called with new opCtx. Ensure that 'child' + * TaskExecutorCursors created from the original TEC's multi-cursor-response can safely + * operate if this happens/don't try and use the now-destroyed operation context. + * See SERVER-69702 for context + */ + void ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructedTest() { + auto lsid = makeLogicalSessionIdForTest(); + opCtx->setLogicalSessionId(lsid); + const auto aggCmd = BSON("aggregate" + << "test" + << "pipeline" + << BSON_ARRAY(BSON("returnMultipleCursors" << true))); + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", aggCmd, opCtx.get()); + TaskExecutorCursor tec = makeTec(rcr); + auto expected = BSON("aggregate" + << "test" + << "pipeline" << BSON_ARRAY(BSON("returnMultipleCursors" << true)) + << "lsid" << lsid.toBSON()); + ASSERT_BSONOBJ_EQ(expected, + scheduleSuccessfulMultiCursorResponse("firstBatch", 1, 2, {0, 0})); + // Before calling getNext (and therefore spawning child TECs), destroy the opCtx + // we used to send the initial query and make a new one. + opCtx.reset(); + opCtx = client->makeOperationContext(); + opCtx->setLogicalSessionId(lsid); + // Use the new opCtx to call getNext. The child TECs should not attempt to read from the + // now dead original opCtx. + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + ASSERT_FALSE(tec.getNext(opCtx.get())); + + auto cursorVec = tec.releaseAdditionalCursors(); + ASSERT_EQUALS(cursorVec.size(), 1); + auto secondCursor = std::move(cursorVec[0]); + + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 2); + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 4); + ASSERT_FALSE(hasReadyRequests()); + + ASSERT_FALSE(secondCursor.getNext(opCtx.get())); + } + + void MultipleCursorsGetMoreWorksTest() { + const auto aggCmd = BSON("aggregate" + << "test" + << "pipeline" + << BSON_ARRAY(BSON("returnMultipleCursors" << true))); + + std::vector<size_t> cursorIds{1, 2}; + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", aggCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr); + + ASSERT_BSONOBJ_EQ(aggCmd, + scheduleSuccessfulMultiCursorResponse("firstBatch", 1, 2, cursorIds)); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + auto cursorVec = tec.releaseAdditionalCursors(); + ASSERT_EQUALS(cursorVec.size(), 1); + + // If we try to getNext() at this point, we are interruptible and can timeout + ASSERT_THROWS_CODE(opCtx->runWithDeadline(Date_t::now() + Milliseconds(100), + ErrorCodes::ExceededTimeLimit, + [&] { tec.getNext(opCtx.get()); }), + DBException, + ErrorCodes::ExceededTimeLimit); + + // We can pick up after that interruption though + ASSERT_BSONOBJ_EQ(BSON("getMore" << 1LL << "collection" + << "test"), + scheduleSuccessfulCursorResponse("nextBatch", 3, 5, cursorIds[0])); + + // Repeat for second cursor. + auto secondCursor = std::move(cursorVec[0]); + + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 2); + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 4); + + ASSERT_THROWS_CODE(opCtx->runWithDeadline(Date_t::now() + Milliseconds(100), + ErrorCodes::ExceededTimeLimit, + [&] { secondCursor.getNext(opCtx.get()); }), + DBException, + ErrorCodes::ExceededTimeLimit); + + ASSERT_BSONOBJ_EQ(BSON("getMore" << 2LL << "collection" + << "test"), + scheduleSuccessfulCursorResponse("nextBatch", 6, 8, cursorIds[1])); + // Read second batch, then schedule EOF on both cursors. + // Then read final document for each. + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 3); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 4); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 5); + scheduleSuccessfulCursorResponse("nextBatch", 6, 6, 0); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 6); + + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 6); + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 7); + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 8); + scheduleSuccessfulCursorResponse("nextBatch", 12, 12, 0); + ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).value()["x"].Int(), 12); + + // Shouldn't have any more requests, both cursors are closed. + ASSERT_FALSE(hasReadyRequests()); + + ASSERT_FALSE(tec.getNext(opCtx.get())); + ASSERT_FALSE(secondCursor.getNext(opCtx.get())); + } + + /** + * Ensure we work if find fails (and that we receive the error code it failed with) + */ + void FailureInFindTest() { + const auto findCmd = BSON("find" + << "test" + << "batchSize" << 2); + + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr); + + scheduleErrorResponse(Status(ErrorCodes::BadValue, "an error")); + + ASSERT_THROWS_CODE(tec.getNext(opCtx.get()), DBException, ErrorCodes::BadValue); + } + + + /** + * Ensure multiple batches works correctly + */ + void MultipleBatchesWorksTest() { + const auto findCmd = BSON("find" + << "test" + << "batchSize" << 2); + CursorId cursorId = 1; + + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr, [] { + TaskExecutorCursor::Options opts; + opts.batchSize = 3; + return opts; + }()); + + scheduleSuccessfulCursorResponse("firstBatch", 1, 2, cursorId); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + // ASSERT(hasReadyRequests()); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + // If we try to getNext() at this point, we are interruptible and can timeout + ASSERT_THROWS_CODE(opCtx->runWithDeadline(Date_t::now() + Milliseconds(100), + ErrorCodes::ExceededTimeLimit, + [&] { tec.getNext(opCtx.get()); }), + DBException, + ErrorCodes::ExceededTimeLimit); + + // We can pick up after that interruption though + ASSERT_BSONOBJ_EQ(BSON("getMore" << 1LL << "collection" + << "test" + << "batchSize" << 3), + scheduleSuccessfulCursorResponse("nextBatch", 3, 5, cursorId)); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 3); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 4); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 5); + + cursorId = 0; + scheduleSuccessfulCursorResponse("nextBatch", 6, 6, cursorId); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 6); + + // We don't issue extra getmores after returning a 0 cursor id + ASSERT_FALSE(hasReadyRequests()); + + ASSERT_FALSE(tec.getNext(opCtx.get())); + } + + /** + * Ensure we allow empty firstBatch. + */ + void EmptyFirstBatchTest() { + const auto findCmd = BSON("find" + << "test" + << "batchSize" << 2); + const auto getMoreCmd = BSON("getMore" << 1LL << "collection" + << "test" + << "batchSize" << 3); + const CursorId cursorId = 1; + + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr, [] { + TaskExecutorCursor::Options opts; + opts.batchSize = 3; + return opts; + }()); + + // Schedule a cursor response with an empty "firstBatch". Use end < start so we don't + // append any doc to "firstBatch". + ASSERT_BSONOBJ_EQ(findCmd, scheduleSuccessfulCursorResponse("firstBatch", 1, 0, cursorId)); + + stdx::thread th([&] { + // Wait for the getMore run by the getNext() below to be ready, and schedule a + // cursor response with a non-empty "nextBatch". + while (!hasReadyRequests()) { + sleepmillis(10); } - cursor.append("id", (long long)(cursorId)); - cursor.append("ns", "test.test"); - } - bob.append("ok", int(1)); + + ASSERT_BSONOBJ_EQ(getMoreCmd, scheduleSuccessfulCursorResponse("nextBatch", 1, 1, 0)); + }); + + // Verify that the first doc is the doc from the second batch. + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + th.join(); + } + + /** + * Ensure we allow any empty non-initial batch. + */ + void EmptyNonInitialBatchTest() { + const auto findCmd = BSON("find" + << "test" + << "batchSize" << 2); + const auto getMoreCmd = BSON("getMore" << 1LL << "collection" + << "test" + << "batchSize" << 3); + const CursorId cursorId = 1; + + RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); + + TaskExecutorCursor tec = makeTec(rcr, [] { + TaskExecutorCursor::Options opts; + opts.batchSize = 3; + return opts; + }()); + + // Schedule a cursor response with a non-empty "firstBatch". + ASSERT_BSONOBJ_EQ(findCmd, scheduleSuccessfulCursorResponse("firstBatch", 1, 1, cursorId)); + + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + + // Schedule two consecutive cursor responses with empty "nextBatch". Use end < start so + // we don't append any doc to "nextBatch". + ASSERT_BSONOBJ_EQ(getMoreCmd, + scheduleSuccessfulCursorResponse("nextBatch", 1, 0, cursorId)); + + stdx::thread th([&] { + // Wait for the first getMore run by the getNext() below to be ready, and schedule a + // cursor response with a non-empty "nextBatch". + while (!hasReadyRequests()) { + sleepmillis(10); + } + + ASSERT_BSONOBJ_EQ(getMoreCmd, + scheduleSuccessfulCursorResponse("nextBatch", 1, 0, cursorId)); + + // Wait for the second getMore run by the getNext() below to be ready, and schedule a + // cursor response with a non-empty "nextBatch". + while (!hasReadyRequests()) { + sleepmillis(10); + } + + ASSERT_BSONOBJ_EQ(getMoreCmd, scheduleSuccessfulCursorResponse("nextBatch", 2, 2, 0)); + }); + + // Verify that the next doc is the doc from the fourth batch. + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + th.join(); + } + + /** + * Test that if 'preFetchNextBatch' is false, the TaskExecutorCursor does not request GetMores + * until the current batch is exhausted and 'getNext()' is invoked. + */ + void NoPrefetchGetMore() { + unittest::threadAssertionMonitoredTest([&](auto& monitor) { + CursorId cursorId = 1; + RemoteCommandRequest rcr(HostAndPort("localhost"), + "test", + BSON("search" + << "foo"), + opCtx.get()); + + // The lambda that will be used to augment the getMore request sent below is passed into + // the TEC constructor. + auto augmentGetMore = [](BSONObjBuilder& bob) { bob.append("test", 1); }; + + // Construction of the TaskExecutorCursor enqueues a request in the + // NetworkInterfaceMock. + TaskExecutorCursor tec = makeTec(rcr, [&augmentGetMore] { + TaskExecutorCursor::Options opts; + opts.batchSize = 2; + opts.preFetchNextBatch = false; + opts.getMoreAugmentationWriter = augmentGetMore; + return opts; + }()); + + // Mock the response for the first batch. + scheduleSuccessfulCursorResponse("firstBatch", 1, 2, cursorId); + + // Exhaust the first batch. + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 1); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 2); + + // Assert that the TaskExecutorCursor has not requested a GetMore. This enforces that + // 'preFetchNextBatch' works as expected. + ASSERT_FALSE(hasReadyRequests()); + + // As soon as 'getNext()' is invoked, the TaskExecutorCursor will try to send a GetMore + // and that will block this thread in the NetworkInterfaceMock until there is a + // scheduled response. However, we cannot schedule the cursor response on the main + // thread before we call 'getNext()' as that will cause the NetworkInterfaceMock to + // block until there is request enqueued ('getNext()' is the function which will enqueue + // such as request). To avoid this deadlock, we start a new thread which will schedule a + // response on the NetworkInterfaceMock. + auto responseSchedulerThread = monitor.spawn([&] { + auto recievedGetMoreCmd = scheduleSuccessfulCursorResponse("nextBatch", 3, 4, 0); + + // Assert that the command processed for the above response matches with the + // lambda to augment the getMore command used during construction of the TEC + // above. + const auto expectedGetMoreCmd = BSON("getMore" << 1LL << "collection" + << "test" + << "batchSize" << 2 << "test" << 1); + ASSERT_BSONOBJ_EQ(expectedGetMoreCmd, recievedGetMoreCmd); + }); + + // Schedules the GetMore request and exhausts the cursor. + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 3); + ASSERT_EQUALS(tec.getNext(opCtx.get()).value()["x"].Int(), 4); + + // Joining the thread which schedules the cursor response for the GetMore here forces + // the destructor of NetworkInterfaceMock::InNetworkGuard to run, which ensures that the + // 'NetworkInterfaceMock' stops executing as the network thread. This is required before + // we invoke 'hasReadyRequests()' which enters the network again. + responseSchedulerThread.join(); + + // Assert no GetMore is requested. + ASSERT_FALSE(hasReadyRequests()); + }); + } + + ServiceContext::UniqueServiceContext serviceCtx = ServiceContext::make(); + ServiceContext::UniqueClient client; + ServiceContext::UniqueOperationContext opCtx; +}; + +class NonPinningTaskExecutorCursorTestFixture + : public TaskExecutorCursorTestFixture<NonPinningTaskExecutorCursorTestFixture, + ThreadPoolExecutorTest> { +public: + void postSetUp() { + launchExecutorThread(); + } + + virtual BSONObj scheduleSuccessfulCursorResponse(StringData fieldName, + size_t start, + size_t end, + size_t cursorId) { + NetworkInterfaceMock::InNetworkGuard ing(getNet()); + ASSERT(getNet()->hasReadyRequests()); - auto rcr = getNet()->scheduleSuccessfulResponse(bob.obj()); + auto rcr = getNet()->scheduleSuccessfulResponse( + buildCursorResponse(fieldName, start, end, cursorId)); getNet()->runReadyNetworkOperations(); return rcr.cmdObj.getOwned(); @@ -102,30 +601,10 @@ public: std::vector<size_t> cursorIds) { NetworkInterfaceMock::InNetworkGuard ing(getNet()); - BSONObjBuilder bob; - { - BSONArrayBuilder cursors; - int baseCursorValue = 1; - for (auto cursorId : cursorIds) { - BSONObjBuilder cursor; - BSONArrayBuilder batch; - ASSERT(start < end && end < INT_MAX); - for (size_t i = start; i <= end; ++i) { - batch.append(BSON("x" << static_cast<int>(i) * baseCursorValue).getOwned()); - } - cursor.append(fieldName, batch.arr()); - cursor.append("id", (long long)(cursorId)); - cursor.append("ns", "test.test"); - auto cursorObj = BSON("cursor" << cursor.done() << "ok" << 1); - cursors.append(cursorObj.getOwned()); - ++baseCursorValue; - } - bob.append("cursors", cursors.arr()); - } - bob.append("ok", 1); ASSERT(getNet()->hasReadyRequests()); - auto rcr = getNet()->scheduleSuccessfulResponse(bob.obj()); + auto rcr = getNet()->scheduleSuccessfulResponse( + buildMultiCursorResponse(fieldName, start, end, cursorIds)); getNet()->runReadyNetworkOperations(); return rcr.cmdObj.getOwned(); @@ -144,167 +623,235 @@ public: return rcr.cmdObj.getOwned(); } + void scheduleErrorResponse(Status error) { + NetworkInterfaceMock::InNetworkGuard ing(getNet()); + + ASSERT(getNet()->hasReadyRequests()); + getNet()->scheduleErrorResponse(error); + getNet()->runReadyNetworkOperations(); + } + bool hasReadyRequests() { NetworkInterfaceMock::InNetworkGuard ing(getNet()); return getNet()->hasReadyRequests(); } - ServiceContext::UniqueServiceContext serviceCtx = ServiceContext::make(); - ServiceContext::UniqueClient client; - ServiceContext::UniqueOperationContext opCtx; -}; - -/** - * Ensure we work for a single simple batch - */ -TEST_F(TaskExecutorCursorFixture, SingleBatchWorks) { - const auto findCmd = BSON("find" - << "test" - << "batchSize" << 2); - const CursorId cursorId = 0; + void blackHoleNextOutgoingRequest() { + NetworkInterfaceMock::InNetworkGuard guard(getNet()); + getNet()->blackHole(getNet()->getFrontOfUnscheduledQueue()); + } - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); + TaskExecutorCursor makeTec(RemoteCommandRequest rcr, + TaskExecutorCursor::Options&& options = {}) { + options.pinConnection = false; + return TaskExecutorCursor(getExecutorPtr(), rcr, std::move(options)); + } +}; - TaskExecutorCursor tec(&getExecutor(), rcr); +class PinnedConnTaskExecutorCursorTestFixture + : public TaskExecutorCursorTestFixture<PinnedConnTaskExecutorCursorTestFixture, + PinnedConnectionTaskExecutorTest> { +public: + void postSetUp() {} + + BSONObj scheduleResponse(StatusWith<BSONObj> response) { + int32_t responseToId; + BSONObj cmdObjReceived; + auto pf = makePromiseFuture<void>(); + expectSinkMessage([&](Message m) { + responseToId = m.header().getId(); + auto opMsg = OpMsgRequest::parse(m); + cmdObjReceived = opMsg.body.removeField("$db").getOwned(); + pf.promise.emplaceValue(); + return Status::OK(); + }); + // Wait until we recieved the command request. + pf.future.get(); + + // Now we expect source message to be called and provide the response + expectSourceMessage([=]() { + rpc::OpMsgReplyBuilder replyBuilder; + replyBuilder.setCommandReply(response); + auto message = replyBuilder.done(); + message.header().setResponseToMsgId(responseToId); + return message; + }); + return cmdObjReceived; + } - ASSERT_BSONOBJ_EQ(findCmd, scheduleSuccessfulCursorResponse("firstBatch", 1, 2, cursorId)); + BSONObj scheduleSuccessfulCursorResponse(StringData fieldName, + size_t start, + size_t end, + size_t cursorId) { + auto cursorResponse = buildCursorResponse(fieldName, start, end, cursorId); + return scheduleResponse(cursorResponse); + } - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 1); + BSONObj scheduleSuccessfulMultiCursorResponse(StringData fieldName, + size_t start, + size_t end, + std::vector<size_t> cursorIds) { + auto cursorResponse = buildMultiCursorResponse(fieldName, start, end, cursorIds); + return scheduleResponse(cursorResponse); + } - ASSERT_FALSE(hasReadyRequests()); + void scheduleErrorResponse(Status error) { + scheduleResponse(error); + } - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 2); + BSONObj scheduleSuccessfulKillCursorResponse(size_t cursorId) { - ASSERT_FALSE(tec.getNext(opCtx.get())); -} + auto cursorResponse = + BSON("cursorsKilled" << BSON_ARRAY((long long)(cursorId)) << "cursorsNotFound" + << BSONArray() << "cursorsAlive" << BSONArray() << "cursorsUnknown" + << BSONArray() << "ok" << 1); + return scheduleResponse(cursorResponse); + } -/** - * Ensure the firstBatch can be read correctly when multiple cursors are returned. - */ -TEST_F(TaskExecutorCursorFixture, MultipleCursorsSingleBatchSucceeds) { - const auto aggCmd = BSON("aggregate" - << "test" - << "pipeline" << BSON_ARRAY(BSON("returnMultipleCursors" << true))); + TaskExecutorCursor makeTec(RemoteCommandRequest rcr, + TaskExecutorCursor::Options&& options = {}) { + options.pinConnection = true; + return TaskExecutorCursor(getExecutorPtr(), rcr, std::move(options)); + } - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", aggCmd, opCtx.get()); + bool hasReadyRequests() { + return asBase().hasReadyRequests(); + } - TaskExecutorCursor tec(&getExecutor(), rcr); + void blackHoleNextOutgoingRequest() { + auto pf = makePromiseFuture<void>(); + expectSinkMessage([&](Message m) { + pf.promise.emplaceValue(); + return Status(ErrorCodes::SocketException, "test"); + }); + pf.future.get(); + } +}; - ASSERT_BSONOBJ_EQ(aggCmd, scheduleSuccessfulMultiCursorResponse("firstBatch", 1, 2, {0, 0})); +class NoPrefetchTaskExecutorCursorTestFixture : public NonPinningTaskExecutorCursorTestFixture { +public: + TaskExecutorCursor makeTec(RemoteCommandRequest rcr, + TaskExecutorCursor::Options&& options = {}) { + options.preFetchNextBatch = false; + return TaskExecutorCursor(getExecutorPtr(), rcr, std::move(options)); + } - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 1); + BSONObj scheduleSuccessfulCursorResponse(StringData fieldName, + size_t start, + size_t end, + size_t cursorId) { + NetworkInterfaceMock::InNetworkGuard ing(getNet()); + // Don't assert that the network has requests like we do in other classes. This is to enable + // the test in 'NoPrefetchGetMore'. + auto rcr = + ing->scheduleSuccessfulResponse(buildCursorResponse(fieldName, start, end, cursorId)); + ing->runReadyNetworkOperations(); + return rcr.cmdObj.getOwned(); + } +}; - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 2); +class NoPrefetchPinnedTaskExecutorCursorTestFixture + : public PinnedConnTaskExecutorCursorTestFixture { +public: + TaskExecutorCursor makeTec(RemoteCommandRequest rcr, + TaskExecutorCursor::Options&& options = {}) { + options.preFetchNextBatch = false; + options.pinConnection = true; + return TaskExecutorCursor(getExecutorPtr(), rcr, std::move(options)); + } +}; - ASSERT_FALSE(tec.getNext(opCtx.get())); +TEST_F(NonPinningTaskExecutorCursorTestFixture, SingleBatchWorks) { + SingleBatchWorksTest(); +} - auto cursorVec = tec.releaseAdditionalCursors(); - ASSERT_EQUALS(cursorVec.size(), 1); - auto secondCursor = std::move(cursorVec[0]); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, SingleBatchWorks) { + SingleBatchWorksTest(); +} - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 2); - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 4); - ASSERT_FALSE(hasReadyRequests()); +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, SingleBatchWorks) { + SingleBatchWorksTest(); +} - ASSERT_FALSE(secondCursor.getNext(opCtx.get())); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, SingleBatchWorks) { + SingleBatchWorksTest(); } -TEST_F(TaskExecutorCursorFixture, MultipleCursorsGetMoreWorks) { - const auto aggCmd = BSON("aggregate" - << "test" - << "pipeline" << BSON_ARRAY(BSON("returnMultipleCursors" << true))); +TEST_F(NonPinningTaskExecutorCursorTestFixture, MultipleCursorsSingleBatchSucceeds) { + MultipleCursorsSingleBatchSucceedsTest(); +} - std::vector<size_t> cursorIds{1, 2}; - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", aggCmd, opCtx.get()); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, MultipleCursorsSingleBatchSucceeds) { + MultipleCursorsSingleBatchSucceedsTest(); +} - TaskExecutorCursor tec(&getExecutor(), rcr); +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, MultipleCursorsSingleBatchSucceeds) { + MultipleCursorsSingleBatchSucceedsTest(); +} - ASSERT_BSONOBJ_EQ(aggCmd, scheduleSuccessfulMultiCursorResponse("firstBatch", 1, 2, cursorIds)); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, MultipleCursorsSingleBatchSucceeds) { + MultipleCursorsSingleBatchSucceedsTest(); +} - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 1); +TEST_F(NonPinningTaskExecutorCursorTestFixture, + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructed) { + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructedTest(); +} - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 2); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructed) { + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructedTest(); +} - auto cursorVec = tec.releaseAdditionalCursors(); - ASSERT_EQUALS(cursorVec.size(), 1); +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructed) { + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructedTest(); +} - // If we try to getNext() at this point, we are interruptible and can timeout - ASSERT_THROWS_CODE(opCtx->runWithDeadline(Date_t::now() + Milliseconds(100), - ErrorCodes::ExceededTimeLimit, - [&] { tec.getNext(opCtx.get()); }), - DBException, - ErrorCodes::ExceededTimeLimit); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructed) { + ChildTaskExecutorCursorsAreSafeIfOriginalOpCtxDestructedTest(); +} - // We can pick up after that interruption though - ASSERT_BSONOBJ_EQ(BSON("getMore" << 1LL << "collection" - << "test"), - scheduleSuccessfulCursorResponse("nextBatch", 3, 5, cursorIds[0])); - - // Repeat for second cursor. - auto secondCursor = std::move(cursorVec[0]); - - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 2); - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 4); - - ASSERT_THROWS_CODE(opCtx->runWithDeadline(Date_t::now() + Milliseconds(100), - ErrorCodes::ExceededTimeLimit, - [&] { secondCursor.getNext(opCtx.get()); }), - DBException, - ErrorCodes::ExceededTimeLimit); - - ASSERT_BSONOBJ_EQ(BSON("getMore" << 2LL << "collection" - << "test"), - scheduleSuccessfulCursorResponse("nextBatch", 6, 8, cursorIds[1])); - // Read second batch on both cursors. - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 3); - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 4); - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 5); - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 6); - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 7); - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 8); - - // Schedule EOF on both cursors. - scheduleSuccessfulCursorResponse("nextBatch", 6, 6, 0); - scheduleSuccessfulCursorResponse("nextBatch", 12, 12, 0); - - // Read final document. - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 6); - ASSERT_EQUALS(secondCursor.getNext(opCtx.get()).get()["x"].Int(), 12); - - // Shouldn't have any more requests, both cursors are closed. - ASSERT_FALSE(hasReadyRequests()); +TEST_F(NonPinningTaskExecutorCursorTestFixture, MultipleCursorsGetMoreWorks) { + MultipleCursorsGetMoreWorksTest(); +} - ASSERT_FALSE(tec.getNext(opCtx.get())); - ASSERT_FALSE(secondCursor.getNext(opCtx.get())); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, MultipleCursorsGetMoreWorks) { + MultipleCursorsGetMoreWorksTest(); } -/** - * Ensure we work if find fails (and that we receive the error code it failed with) - */ -TEST_F(TaskExecutorCursorFixture, FailureInFind) { - const auto findCmd = BSON("find" - << "test" - << "batchSize" << 2); +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, MultipleCursorsGetMoreWorks) { + MultipleCursorsGetMoreWorksTest(); +} - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, MultipleCursorsGetMoreWorks) { + MultipleCursorsGetMoreWorksTest(); +} - TaskExecutorCursor tec(&getExecutor(), rcr); +TEST_F(NonPinningTaskExecutorCursorTestFixture, FailureInFind) { + FailureInFindTest(); +} - { - NetworkInterfaceMock::InNetworkGuard ing(getNet()); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, FailureInFind) { + FailureInFindTest(); +} - ASSERT(getNet()->hasReadyRequests()); - getNet()->scheduleErrorResponse(Status(ErrorCodes::BadValue, "an error")); - getNet()->runReadyNetworkOperations(); - } +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, FailureInFind) { + FailureInFindTest(); +} - ASSERT_THROWS_CODE(tec.getNext(opCtx.get()), DBException, ErrorCodes::BadValue); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, FailureInFind) { + FailureInFindTest(); } /** * Ensure early termination of the cursor calls killCursor (if we know about the cursor id) + * Only applicable to the unpinned case - if the connection is pinned, and a getMore is + * in progress and/or fails, the most we can do is kill the connection. We can't re-use + * the connection to send killCursors. */ -TEST_F(TaskExecutorCursorFixture, EarlyReturnKillsCursor) { +TEST_F(NonPinningTaskExecutorCursorTestFixture, EarlyReturnKillsCursor) { const auto findCmd = BSON("find" << "test" << "batchSize" << 2); @@ -313,172 +860,76 @@ TEST_F(TaskExecutorCursorFixture, EarlyReturnKillsCursor) { RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); { - TaskExecutorCursor tec(&getExecutor(), rcr); + TaskExecutorCursor tec = makeTec(rcr); scheduleSuccessfulCursorResponse("firstBatch", 1, 2, cursorId); ASSERT(tec.getNext(opCtx.get())); + + // Black hole the pending `getMore` operation scheduled by the `TaskExecutorCursor`. + blackHoleNextOutgoingRequest(); } + ASSERT_BSONOBJ_EQ(BSON("killCursors" << "test" << "cursors" << BSON_ARRAY(1)), scheduleSuccessfulKillCursorResponse(1)); } -/** - * Ensure multiple batches works correctly - */ -TEST_F(TaskExecutorCursorFixture, MultipleBatchesWorks) { - const auto findCmd = BSON("find" - << "test" - << "batchSize" << 2); - CursorId cursorId = 1; - - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); - - TaskExecutorCursor tec(&getExecutor(), rcr, [] { - TaskExecutorCursor::Options opts; - opts.batchSize = 3; - return opts; - }()); - - scheduleSuccessfulCursorResponse("firstBatch", 1, 2, cursorId); - - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 1); - - ASSERT(hasReadyRequests()); - - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 2); - - // If we try to getNext() at this point, we are interruptible and can timeout - ASSERT_THROWS_CODE(opCtx->runWithDeadline(Date_t::now() + Milliseconds(100), - ErrorCodes::ExceededTimeLimit, - [&] { tec.getNext(opCtx.get()); }), - DBException, - ErrorCodes::ExceededTimeLimit); - - // We can pick up after that interruption though - ASSERT_BSONOBJ_EQ(BSON("getMore" << 1LL << "collection" - << "test" - << "batchSize" << 3), - scheduleSuccessfulCursorResponse("nextBatch", 3, 5, cursorId)); - - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 3); - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 4); - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 5); - - cursorId = 0; - scheduleSuccessfulCursorResponse("nextBatch", 6, 6, cursorId); - - // We don't issue extra getmores after returning a 0 cursor id - ASSERT_FALSE(hasReadyRequests()); - - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 6); - - ASSERT_FALSE(tec.getNext(opCtx.get())); +TEST_F(NonPinningTaskExecutorCursorTestFixture, MultipleBatchesWorks) { + MultipleBatchesWorksTest(); } -/** - * Ensure we allow empty firstBatch. - */ -TEST_F(TaskExecutorCursorFixture, EmptyFirstBatch) { - const auto findCmd = BSON("find" - << "test" - << "batchSize" << 2); - const auto getMoreCmd = BSON("getMore" << 1LL << "collection" - << "test" - << "batchSize" << 3); - const CursorId cursorId = 1; - - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); - - TaskExecutorCursor tec(&getExecutor(), rcr, [] { - TaskExecutorCursor::Options opts; - opts.batchSize = 3; - return opts; - }()); - - // Schedule a cursor response with an empty "firstBatch". Use end < start so we don't - // append any doc to "firstBatch". - ASSERT_BSONOBJ_EQ(findCmd, scheduleSuccessfulCursorResponse("firstBatch", 1, 0, cursorId)); - - stdx::thread th([&] { - // Wait for the getMore run by the getNext() below to be ready, and schedule a - // cursor response with a non-empty "nextBatch". - while (!hasReadyRequests()) { - sleepmillis(10); - } - - ASSERT_BSONOBJ_EQ(getMoreCmd, - scheduleSuccessfulCursorResponse("nextBatch", 1, 1, cursorId)); - }); - - // Verify that the first doc is the doc from the second batch. - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 1); - - th.join(); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, MultipleBatchesWorks) { + MultipleBatchesWorksTest(); } -/** - * Ensure we allow any empty non-initial batch. - */ -TEST_F(TaskExecutorCursorFixture, EmptyNonInitialBatch) { - const auto findCmd = BSON("find" - << "test" - << "batchSize" << 2); - const auto getMoreCmd = BSON("getMore" << 1LL << "collection" - << "test" - << "batchSize" << 3); - const CursorId cursorId = 1; - - RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); - - TaskExecutorCursor tec(&getExecutor(), rcr, [] { - TaskExecutorCursor::Options opts; - opts.batchSize = 3; - return opts; - }()); +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, MultipleBatchesWorks) { + MultipleBatchesWorksTest(); +} - // Schedule a cursor response with a non-empty "firstBatch". - ASSERT_BSONOBJ_EQ(findCmd, scheduleSuccessfulCursorResponse("firstBatch", 1, 1, cursorId)); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, MultipleBatchesWorks) { + MultipleBatchesWorksTest(); +} - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 1); +TEST_F(NonPinningTaskExecutorCursorTestFixture, EmptyFirstBatch) { + EmptyFirstBatchTest(); +} - // Schedule two consecutive cursor responses with empty "nextBatch". Use end < start so - // we don't append any doc to "nextBatch". - ASSERT_BSONOBJ_EQ(getMoreCmd, scheduleSuccessfulCursorResponse("nextBatch", 1, 0, cursorId)); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, EmptyFirstBatch) { + EmptyFirstBatchTest(); +} - stdx::thread th([&] { - // Wait for the first getMore run by the getNext() below to be ready, and schedule a - // cursor response with a non-empty "nextBatch". - while (!hasReadyRequests()) { - sleepmillis(10); - } +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, EmptyFirstBatch) { + EmptyFirstBatchTest(); +} - ASSERT_BSONOBJ_EQ(getMoreCmd, - scheduleSuccessfulCursorResponse("nextBatch", 1, 0, cursorId)); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, EmptyFirstBatch) { + EmptyFirstBatchTest(); +} - // Wait for the second getMore run by the getNext() below to be ready, and schedule a - // cursor response with a non-empty "nextBatch". - while (!hasReadyRequests()) { - sleepmillis(10); - } +TEST_F(NonPinningTaskExecutorCursorTestFixture, EmptyNonInitialBatch) { + EmptyNonInitialBatchTest(); +} - ASSERT_BSONOBJ_EQ(getMoreCmd, - scheduleSuccessfulCursorResponse("nextBatch", 2, 2, cursorId)); - }); +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, EmptyNonInitialBatch) { + EmptyNonInitialBatchTest(); +} - // Verify that the next doc is the doc from the fourth batch. - ASSERT_EQUALS(tec.getNext(opCtx.get()).get()["x"].Int(), 2); +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, EmptyNonInitialBatch) { + EmptyNonInitialBatchTest(); +} - th.join(); +TEST_F(PinnedConnTaskExecutorCursorTestFixture, EmptyNonInitialBatch) { + EmptyNonInitialBatchTest(); } /** - * Ensure lsid is passed in all stages of querying + * Ensure the LSID is passed in all stages of querying. Need to test the + * pinning case separately because of difference around killCursor. */ -TEST_F(TaskExecutorCursorFixture, LsidIsPassed) { +TEST_F(NonPinningTaskExecutorCursorTestFixture, LsidIsPassed) { auto lsid = makeLogicalSessionIdForTest(); opCtx->setLogicalSessionId(lsid); @@ -490,11 +941,11 @@ TEST_F(TaskExecutorCursorFixture, LsidIsPassed) { RemoteCommandRequest rcr(HostAndPort("localhost"), "test", findCmd, opCtx.get()); boost::optional<TaskExecutorCursor> tec; - tec.emplace(&getExecutor(), rcr, []() { + tec.emplace(makeTec(rcr, []() { TaskExecutorCursor::Options opts; opts.batchSize = 1; return opts; - }()); + }())); // lsid in the first batch ASSERT_BSONOBJ_EQ(BSON("find" @@ -521,6 +972,14 @@ TEST_F(TaskExecutorCursorFixture, LsidIsPassed) { ASSERT_FALSE(hasReadyRequests()); } +TEST_F(NoPrefetchTaskExecutorCursorTestFixture, NoPrefetchGetMore) { + NoPrefetchGetMore(); +} + +TEST_F(NoPrefetchPinnedTaskExecutorCursorTestFixture, NoPrefetchWithPinning) { + NoPrefetchGetMore(); +} + } // namespace } // namespace executor } // namespace mongo diff --git a/src/mongo/executor/test_network_connection_hook.h b/src/mongo/executor/test_network_connection_hook.h index 5a22ac24f54..c825903c7dc 100644 --- a/src/mongo/executor/test_network_connection_hook.h +++ b/src/mongo/executor/test_network_connection_hook.h @@ -53,8 +53,8 @@ public: Status validateHost(const HostAndPort& remoteHost, const BSONObj& request, - const RemoteCommandResponse& isMasterReply) override { - return _validateFunc(remoteHost, request, isMasterReply); + const RemoteCommandResponse& helloReply) override { + return _validateFunc(remoteHost, request, helloReply); } StatusWith<boost::optional<RemoteCommandRequest>> makeRequest(const HostAndPort& remoteHost) { diff --git a/src/mongo/executor/thread_pool_task_executor.cpp b/src/mongo/executor/thread_pool_task_executor.cpp index d642252c7f5..c61db6f38b4 100644 --- a/src/mongo/executor/thread_pool_task_executor.cpp +++ b/src/mongo/executor/thread_pool_task_executor.cpp @@ -88,8 +88,9 @@ public: } // All fields except for "canceled" are guarded by the owning task executor's _mutex. The - // "canceled" field may be observed without holding _mutex, but may only be set while holding - // _mutex. + // "canceled" field may be observed without holding _mutex only if we are checking if the value + // is true. This is because once "canceled" stores true, we never set it back to false. The + // "canceled" field may only be set while holding _mutex. CallbackFn callback; AtomicWord<unsigned> canceled{0U}; @@ -823,18 +824,18 @@ void ThreadPoolTaskExecutor::runCallbackExhaust(std::shared_ptr<CallbackState> c std::move(cbHandle), cbState->canceled.load() ? kCallbackCanceledErrorStatus : Status::OK()); - if (!cbState->isFinished.load()) { + if (auto lk = stdx::unique_lock(_mutex); !cbState->isFinished.load()) { TaskExecutor::CallbackFn callback = [](const CallbackArgs&) {}; { - auto lk = stdx::lock_guard(_mutex); std::swap(cbState->callback, callback); + lk.unlock(); } callback(std::move(args)); + lk.lock(); // Leave the empty callback function if the request has been marked canceled or finished // while running the callback to avoid leaking resources. if (!cbState->canceled.load() && !cbState->isFinished.load()) { - auto lk = stdx::lock_guard(_mutex); std::swap(callback, cbState->callback); } } diff --git a/src/mongo/executor/thread_pool_task_executor.h b/src/mongo/executor/thread_pool_task_executor.h index 9c95467118f..900ed7508c8 100644 --- a/src/mongo/executor/thread_pool_task_executor.h +++ b/src/mongo/executor/thread_pool_task_executor.h @@ -231,6 +231,9 @@ private: // Lifecycle state of this executor. stdx::condition_variable _stateChange; State _state = preStart; + + friend std::shared_ptr<TaskExecutor> makePinnedConnectionTaskExecutor( + std::shared_ptr<TaskExecutor>); }; } // namespace executor |