FullTableScanner.hpp

#pragma once
 
#include "etl/ETLHelpers.hpp"
#include "util/Assert.hpp"
#include "util/async/AnyExecutionContext.hpp"
#include "util/async/AnyOperation.hpp"
#include "util/async/context/BasicExecutionContext.hpp"
 
#include <boost/asio/spawn.hpp>
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <ranges>
#include <vector>
 
namespace migration::cassandra::impl {

struct TokenRange {
    std::int64_t start;
    std::int64_t end;

    TokenRange(std::int64_t start, std::int64_t end) : start{start}, end{end}
    {
    }
};

template <typename T>
concept CanReadByTokenRange =
    requires(T obj, TokenRange const& range, boost::asio::yield_context yield) {
        { obj.readByTokenRange(range, yield) } -> std::same_as<void>;
    };

template <CanReadByTokenRange TableAdapter>
class FullTableScanner {
    struct TokenRangesProvider {
        uint32_t numRanges;
 
        TokenRangesProvider(uint32_t numRanges) : numRanges{numRanges}
        {
        }
 
        [[nodiscard]] std::vector<TokenRange>
        getRanges() const
        {
            auto const minValue = std::numeric_limits<std::int64_t>::min();
            auto const maxValue = std::numeric_limits<std::int64_t>::max();
            if (numRanges == 1)
                return {TokenRange{minValue, maxValue}};
 
            // Safely calculate the range size using uint64_t to avoid overflow
            uint64_t const rangeSize = (static_cast<uint64_t>(maxValue) * 2) / numRanges;
 
            std::vector<TokenRange> ranges;
            ranges.reserve(numRanges);
 
            for (std::int64_t i = 0; i < numRanges; ++i) {
                int64_t const start = minValue + (i * rangeSize);
                int64_t const end =
                    (i == numRanges - 1) ? maxValue : start + static_cast<int64_t>(rangeSize) - 1;
                ranges.emplace_back(start, end);
            }
 
            return ranges;
        }
    };
 
    [[nodiscard]] auto
    spawnWorker()
    {
        return ctx_.execute([this](auto token) {
            while (not token.isStopRequested()) {
                auto cursor = queue_.tryPop();
                if (not cursor.has_value()) {
                    return;  // queue is empty
                }
                reader_.readByTokenRange(cursor.value(), token);
            }
        });
    }
 
    void
    load(size_t workerNum)
    {
        namespace vs = std::views;
 
        tasks_.reserve(workerNum);
 
        for ([[maybe_unused]] auto const taskId : vs::iota(0u, workerNum))
            tasks_.push_back(spawnWorker());
    }
 
    util::async::AnyExecutionContext ctx_;
    std::size_t cursorsNum_;
    etl::ThreadSafeQueue<TokenRange> queue_;
    std::vector<util::async::AnyOperation<void>> tasks_;
    TableAdapter reader_;
 
public:
    struct FullTableScannerSettings {
        std::uint32_t ctxThreadsNum;  
        std::uint32_t jobsNum;  
        std::uint32_t cursorsPerJob;  
    };

    template <typename ExecutionContextType = util::async::CoroExecutionContext>
    FullTableScanner(FullTableScannerSettings settings, TableAdapter&& reader)
        : ctx_(ExecutionContextType(settings.ctxThreadsNum))
        , cursorsNum_(settings.jobsNum * settings.cursorsPerJob)
        , queue_{cursorsNum_}
        , reader_{std::move(reader)}
    {
        ASSERT(settings.jobsNum > 0, "jobsNum for full table scanner must be greater than 0");
        ASSERT(
            settings.cursorsPerJob > 0,
            "cursorsPerJob for full table scanner must be greater than 0"
        );
 
        auto const cursors = TokenRangesProvider{cursorsNum_}.getRanges();
        std::ranges::for_each(cursors, [this](auto const& cursor) { queue_.push(cursor); });
        load(settings.jobsNum);
    }

    void
    wait()
    {
        for (auto& task : tasks_) {
            task.wait();
        }
    }
};
 
}  // namespace migration::cassandra::impl