ExecutionStrategy.hpp

//------------------------------------------------------------------------------
/*
    This file is part of clio: https://github.com/XRPLF/clio
    Copyright (c) 2023, the clio developers.
 
    Permission to use, copy, modify, and distribute this software for any
    purpose with or without fee is hereby granted, provided that the above
    copyright notice and this permission notice appear in all copies.
 
    THE  SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
    WITH  REGARD  TO  THIS  SOFTWARE  INCLUDING  ALL  IMPLIED  WARRANTIES  OF
    MERCHANTABILITY  AND  FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
    ANY  SPECIAL,  DIRECT,  INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
    WHATSOEVER  RESULTING  FROM  LOSS  OF USE, DATA OR PROFITS, WHETHER IN AN
    ACTION  OF  CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
*/
//==============================================================================
 
#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