ExecutionStrategy.hpp

#pragma once
 
#include "data/BackendCounters.hpp"
#include "data/BackendInterface.hpp"
#include "data/cassandra/Handle.hpp"
#include "data/cassandra/Types.hpp"
#include "data/cassandra/impl/AsyncExecutor.hpp"
#include "util/Assert.hpp"
#include "util/Batching.hpp"
#include "util/log/Logger.hpp"
 
#include <boost/asio.hpp>
#include <boost/asio/associated_executor.hpp>
#include <boost/asio/executor_work_guard.hpp>
#include <boost/asio/io_context.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/json/object.hpp>
 
#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <thread>
#include <type_traits>
#include <vector>
 
namespace data::cassandra::impl {
 
// TODO: this could probably be also moved out of impl and into the main cassandra namespace.

template <typename HandleType = Handle, SomeBackendCounters BackendCountersType = BackendCounters>
class DefaultExecutionStrategy {
    util::Logger log_{"Backend"};
 
    std::uint32_t maxWriteRequestsOutstanding_;
    std::atomic_uint32_t numWriteRequestsOutstanding_ = 0;
 
    std::uint32_t maxReadRequestsOutstanding_;
    std::atomic_uint32_t numReadRequestsOutstanding_ = 0;
 
    std::size_t writeBatchSize_;
 
    std::mutex throttleMutex_;
    std::condition_variable throttleCv_;
 
    std::mutex syncMutex_;
    std::condition_variable syncCv_;
 
    boost::asio::io_context ioc_;
    std::optional<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>> work_;
 
    std::reference_wrapper<HandleType const> handle_;
    std::thread thread_;
 
    typename BackendCountersType::PtrType counters_;
 
public:
    using ResultOrErrorType = typename HandleType::ResultOrErrorType;
    using StatementType = typename HandleType::StatementType;
    using PreparedStatementType = typename HandleType::PreparedStatementType;
    using FutureType = typename HandleType::FutureType;
    using FutureWithCallbackType = typename HandleType::FutureWithCallbackType;
    using ResultType = typename HandleType::ResultType;
    using CompletionTokenType = boost::asio::yield_context;

    DefaultExecutionStrategy(
        Settings const& settings,
        HandleType const& handle,
        typename BackendCountersType::PtrType counters = BackendCountersType::make()
    )
        : maxWriteRequestsOutstanding_{settings.maxWriteRequestsOutstanding}
        , maxReadRequestsOutstanding_{settings.maxReadRequestsOutstanding}
        , writeBatchSize_{settings.writeBatchSize}
        , work_{boost::asio::make_work_guard(ioc_)}
        , handle_{std::cref(handle)}
        , thread_{[this]() { ioc_.run(); }}
        , counters_{std::move(counters)}
    {
        LOG(log_.info()) << "Max write requests outstanding is " << maxWriteRequestsOutstanding_
                         << "; Max read requests outstanding is " << maxReadRequestsOutstanding_;
    }
 
    ~DefaultExecutionStrategy()
    {
        work_.reset();
        ioc_.stop();
        thread_.join();
    }

    void
    sync()
    {
        LOG(log_.debug()) << "Waiting to sync all writes...";
        std::unique_lock<std::mutex> lck(syncMutex_);
        syncCv_.wait(lck, [this]() { return finishedAllWriteRequests(); });
        LOG(log_.debug()) << "Sync done.";
    }

    bool
    isTooBusy() const
    {
        bool const result = numReadRequestsOutstanding_ >= maxReadRequestsOutstanding_;
        if (result)
            counters_->registerTooBusy();
        return result;
    }

    ResultOrErrorType
    writeSync(StatementType const& statement)
    {
        auto const startTime = std::chrono::steady_clock::now();
        while (true) {
            auto res = handle_.get().execute(statement);
            if (res) {
                counters_->registerWriteSync(startTime);
                return res;
            }
 
            counters_->registerWriteSyncRetry();
            LOG(log_.warn()) << "Cassandra sync write error, retrying: " << res.error();
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
        }
    }

    template <typename... Args>
    ResultOrErrorType
    writeSync(PreparedStatementType const& preparedStatement, Args&&... args)
    {
        return writeSync(preparedStatement.bind(std::forward<Args>(args)...));
    }

    template <typename... Args>
    void
    write(PreparedStatementType const& preparedStatement, Args&&... args)
    {
        auto statement = preparedStatement.bind(std::forward<Args>(args)...);
        write(std::move(statement));
    }

    void
    write(StatementType&& statement)
    {
        auto const startTime = std::chrono::steady_clock::now();
 
        incrementOutstandingRequestCount();
 
        counters_->registerWriteStarted();
        // Note: lifetime is controlled by std::shared_from_this internally
        AsyncExecutor<std::decay_t<decltype(statement)>, HandleType>::run(
            ioc_,
            handle_,
            std::move(statement),
            [this, startTime](auto const&) {
                decrementOutstandingRequestCount();
 
                counters_->registerWriteFinished(startTime);
            },
            [this]() { counters_->registerWriteRetry(); }
        );
    }

    void
    write(std::vector<StatementType>&& statements)
    {
        if (statements.empty())
            return;
 
        util::forEachBatch(std::move(statements), writeBatchSize_, [this](auto begin, auto end) {
            auto const startTime = std::chrono::steady_clock::now();
            auto chunk = std::vector<StatementType>{};
 
            chunk.reserve(std::distance(begin, end));
            std::move(begin, end, std::back_inserter(chunk));
 
            incrementOutstandingRequestCount();
            counters_->registerWriteStarted();
 
            // Note: lifetime is controlled by std::shared_from_this internally
            AsyncExecutor<std::decay_t<decltype(chunk)>, HandleType>::run(
                ioc_,
                handle_,
                std::move(chunk),
                [this, startTime](auto const&) {
                    decrementOutstandingRequestCount();
                    counters_->registerWriteFinished(startTime);
                },
                [this]() { counters_->registerWriteRetry(); }
            );
        });
    }

    void
    writeEach(std::vector<StatementType>&& statements)
    {
        std::ranges::for_each(std::move(statements), [this](auto& statement) {
            this->write(std::move(statement));
        });
    }

    template <typename... Args>
    [[maybe_unused]] ResultOrErrorType
    read(CompletionTokenType token, PreparedStatementType const& preparedStatement, Args&&... args)
    {
        return read(token, preparedStatement.bind(std::forward<Args>(args)...));
    }

    [[maybe_unused]] ResultOrErrorType
    read(CompletionTokenType token, std::vector<StatementType> const& statements)
    {
        auto const startTime = std::chrono::steady_clock::now();
 
        auto const numStatements = statements.size();
        std::optional<FutureWithCallbackType> future;
        counters_->registerReadStarted(numStatements);
 
        // todo: perhaps use policy instead
        while (true) {
            numReadRequestsOutstanding_ += numStatements;
 
            auto init = [this, &statements, &future]<typename Self>(Self& self) {
                auto sself = std::make_shared<Self>(std::move(self));
 
                future.emplace(handle_.get().asyncExecute(statements, [sself](auto&& res) mutable {
                    boost::asio::post(
                        boost::asio::get_associated_executor(*sself),
                        [sself, res = std::forward<decltype(res)>(res)]() mutable {
                            sself->complete(std::move(res));
                        }
                    );
                }));
            };
 
            auto res = boost::asio::async_compose<CompletionTokenType, void(ResultOrErrorType)>(
                std::move(init), token, boost::asio::get_associated_executor(token)
            );
            numReadRequestsOutstanding_ -= numStatements;
 
            if (res) {
                counters_->registerReadFinished(startTime, numStatements);
                return res;
            }
 
            LOG(log_.error()) << "Failed batch read in coroutine: " << res.error();
            try {
                throwErrorIfNeeded(res.error());
            } catch (...) {
                counters_->registerReadError(numStatements);
                throw;
            }
            counters_->registerReadRetry(numStatements);
        }
    }

    [[maybe_unused]] ResultOrErrorType
    read(CompletionTokenType token, StatementType const& statement)
    {
        auto const startTime = std::chrono::steady_clock::now();
 
        std::optional<FutureWithCallbackType> future;
        counters_->registerReadStarted();
 
        // todo: perhaps use policy instead
        while (true) {
            ++numReadRequestsOutstanding_;
            auto init = [this, &statement, &future]<typename Self>(Self& self) {
                auto sself = std::make_shared<Self>(std::move(self));
 
                future.emplace(handle_.get().asyncExecute(statement, [sself](auto&& res) mutable {
                    boost::asio::post(
                        boost::asio::get_associated_executor(*sself),
                        [sself, res = std::forward<decltype(res)>(res)]() mutable {
                            sself->complete(std::move(res));
                        }
                    );
                }));
            };
 
            auto res = boost::asio::async_compose<CompletionTokenType, void(ResultOrErrorType)>(
                std::move(init), token, boost::asio::get_associated_executor(token)
            );
            --numReadRequestsOutstanding_;
 
            if (res) {
                counters_->registerReadFinished(startTime);
                return res;
            }
 
            LOG(log_.error()) << "Failed read in coroutine: " << res.error();
            try {
                throwErrorIfNeeded(res.error());
            } catch (...) {
                counters_->registerReadError();
                throw;
            }
            counters_->registerReadRetry();
        }
    }

    std::vector<ResultType>
    readEach(CompletionTokenType token, std::vector<StatementType> const& statements)
    {
        auto const startTime = std::chrono::steady_clock::now();
 
        std::atomic_uint64_t errorsCount = 0u;
        std::atomic_int numOutstanding = statements.size();
        numReadRequestsOutstanding_ += statements.size();
 
        auto futures = std::vector<FutureWithCallbackType>{};
        futures.reserve(numOutstanding);
        counters_->registerReadStarted(statements.size());
 
        auto init = [this, &statements, &futures, &errorsCount, &numOutstanding]<typename Self>(
                        Self& self
                    ) {
            auto sself = std::make_shared<Self>(std::move(self));
            auto executionHandler =
                [&errorsCount, &numOutstanding, sself](auto const& res) mutable {
                    if (not res)
                        ++errorsCount;
 
                    // when all async operations complete unblock the result
                    if (--numOutstanding == 0) {
                        boost::asio::post(
                            boost::asio::get_associated_executor(*sself),
                            [sself]() mutable { sself->complete(); }
                        );
                    }
                };
 
            std::transform(
                std::cbegin(statements),
                std::cend(statements),
                std::back_inserter(futures),
                [this, &executionHandler](auto const& statement) {
                    return handle_.get().asyncExecute(statement, executionHandler);
                }
            );
        };
 
        boost::asio::async_compose<CompletionTokenType, void()>(
            std::move(init), token, boost::asio::get_associated_executor(token)
        );
        numReadRequestsOutstanding_ -= statements.size();
 
        if (errorsCount > 0) {
            ASSERT(
                errorsCount <= statements.size(), "Errors number cannot exceed statements number"
            );
            counters_->registerReadError(errorsCount);
            counters_->registerReadFinished(startTime, statements.size() - errorsCount);
            throw DatabaseTimeout{};
        }
        counters_->registerReadFinished(startTime, statements.size());
 
        std::vector<ResultType> results;
        results.reserve(futures.size());
 
        // it's safe to call blocking get on futures here as we already waited for the coroutine to
        // resume above.
        std::transform(
            std::make_move_iterator(std::begin(futures)),
            std::make_move_iterator(std::end(futures)),
            std::back_inserter(results),
            [](auto&& future) {
                auto entry = future.get();
                auto&& res = entry.value();
                return std::move(res);
            }
        );
 
        ASSERT(
            futures.size() == statements.size(),
            "Futures size must be equal to statements size. Got {} and {}",
            futures.size(),
            statements.size()
        );
        ASSERT(
            results.size() == statements.size(),
            "Results size must be equal to statements size. Got {} and {}",
            results.size(),
            statements.size()
        );
        return results;
    }

    boost::json::object
    stats() const
    {
        return counters_->report();
    }
 
private:
    void
    incrementOutstandingRequestCount()
    {
        {
            std::unique_lock<std::mutex> lck(throttleMutex_);
            if (!canAddWriteRequest()) {
                LOG(log_.trace()) << "Max outstanding requests reached. "
                                  << "Waiting for other requests to finish";
                throttleCv_.wait(lck, [this]() { return canAddWriteRequest(); });
            }
        }
        ++numWriteRequestsOutstanding_;
    }
 
    void
    decrementOutstandingRequestCount()
    {
        // sanity check
        ASSERT(numWriteRequestsOutstanding_ > 0, "Decrementing num outstanding below 0");
        size_t const cur = (--numWriteRequestsOutstanding_);
        {
            // mutex lock required to prevent race condition around spurious
            // wakeup
            std::lock_guard const lck(throttleMutex_);
            throttleCv_.notify_one();
        }
        if (cur == 0) {
            // mutex lock required to prevent race condition around spurious
            // wakeup
            std::lock_guard const lck(syncMutex_);
            syncCv_.notify_one();
        }
    }
 
    bool
    canAddWriteRequest() const
    {
        return numWriteRequestsOutstanding_ < maxWriteRequestsOutstanding_;
    }
 
    bool
    finishedAllWriteRequests() const
    {
        return numWriteRequestsOutstanding_ == 0;
    }
 
    void
    throwErrorIfNeeded(CassandraError err) const
    {
        if (err.isTimeout())
            throw DatabaseTimeout();
 
        if (err.isInvalidQuery())
            throw std::runtime_error("Invalid query");
    }
};
 
}  // namespace data::cassandra::impl