diff options
Diffstat (limited to 'src/mongo/executor/pinned_connection_task_executor.cpp')
| -rw-r--r-- | src/mongo/executor/pinned_connection_task_executor.cpp | 443 |
1 files changed, 443 insertions, 0 deletions
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 |