CherryPickFastCheckpointer.hpp

// <CherryPickFastCheckpointer> -*- C++ -*-
 
#pragma once
 
#include "sparta/serialization/checkpoint/FastCheckpointer.hpp"
#include "sparta/utils/SpartaException.hpp"
#include "simdb/apps/App.hpp"
#include "simdb/utils/ConcurrentQueue.hpp"
#include "simdb/sqlite/Iterator.hpp"
#include <map>
 
namespace sparta
{
    class TreeNode;
    class Scheduler;
}
 
namespace simdb::pipeline
{
    class Flusher;
}
 
namespace sparta::serialization::checkpoint
{
 
class CherryPickFastCheckpointer final : public simdb::App
{
public:
    using checkpoint_type = typename FastCheckpointer::checkpoint_type;
    using checkpoint_ptr = typename FastCheckpointer::checkpoint_ptr;
    using checkpoint_ptrs = typename FastCheckpointer::checkpoint_ptrs;
    using chkpt_id_t = typename Checkpointer::chkpt_id_t;
    using arch_id_t = uint64_t;
    using tick_t = typename Checkpointer::tick_t;
 
    static constexpr auto NAME = "cherry-pick-fast-checkpointer";
 
    CherryPickFastCheckpointer(simdb::DatabaseManager* db_mgr, const std::vector<TreeNode*> & roots,
                               Scheduler* sched = nullptr);
 
    class AppFactory : public simdb::AppFactoryBase
    {
    public:
        using AppT = serialization::checkpoint::CherryPickFastCheckpointer;
 
        void parameterize(const std::vector<TreeNode*>& roots, Scheduler* sched = nullptr)
        {
            roots_ = roots;
            scheduler_ = sched;
        }
 
        AppT* createApp(simdb::DatabaseManager* db_mgr) override
        {
            // Make the ctor call that the default AppFactory cannot make.
            return new AppT(db_mgr, roots_, scheduler_);
        }
 
        void defineSchema(simdb::Schema& schema) const override
        {
            AppT::defineSchema(schema);
        }
 
    private:
        std::vector<TreeNode*> roots_;
        Scheduler* scheduler_ = nullptr;
    };
 
    static void defineSchema(simdb::Schema& schema);
 
    void createPipeline(simdb::pipeline::PipelineManager* pipeline_mgr) override;
 
    FastCheckpointer& getFastCheckpointer() noexcept {
        return checkpointer_;
    }
 
    void commitCurrentBranch(bool force_new_head_chkpt = false,
                             std::vector<chkpt_id_t>* committed_checkpoints = nullptr);
 
    void saveCheckpoints(checkpoint_ptrs&& checkpoints);
 
    void preTeardown() override;
 
    struct ChkptWindow;
 
    struct ArchIdsForTick {
        arch_id_t start_arch_id = UINT64_MAX;
        arch_id_t end_arch_id = UINT64_MAX;
        tick_t tick = UINT64_MAX;
    };
 
    class DatabaseCheckpointLoader
    {
    public:
        DatabaseCheckpointLoader(simdb::DatabaseManager* db_mgr, TreeNode* root) :
            DatabaseCheckpointLoader(db_mgr, std::vector<TreeNode*>({root}))
        {}
 
        DatabaseCheckpointLoader(simdb::DatabaseManager* db_mgr, const std::vector<TreeNode*>& roots);
 
        uint64_t getNumCheckpoints() const { return num_checkpoints_; }
 
        void loadCheckpoint(uint64_t checkpoint_idx) {
            if (checkpoint_idx >= num_checkpoints_) {
                throw SpartaException("Checkpoint index out of range");
            }
            loadCheckpoint_(checkpoint_idx);
        }
 
    private:
        static std::vector<ArchData*> enumerateArchDatas_(const std::vector<TreeNode*>& roots);
 
        static checkpoint_type* getChkptFromWindow_(arch_id_t arch_id, tick_t tick, ChkptWindow& window);
 
        checkpoint_type* getChkptFromCache_(arch_id_t arch_id, tick_t tick);
 
        checkpoint_type* getChkptFromDisk_(arch_id_t arch_id, tick_t tick);
 
        void loadCheckpoint_(uint64_t checkpoint_idx);
 
        simdb::DatabaseManager *const db_mgr_;
 
        const std::vector<ArchData*> adatas_;
 
        std::unique_ptr<ChkptWindow> cached_window_;
 
        std::vector<ArchIdsForTick> tick_runs_;
 
        uint64_t num_checkpoints_ = 0;
    };
 
    class DatabaseCheckpointReplayer
    {
    public:
        DatabaseCheckpointReplayer(simdb::DatabaseManager* db_mgr, TreeNode* root) :
            DatabaseCheckpointReplayer(db_mgr, std::vector<TreeNode*>({root}))
        {}
 
        DatabaseCheckpointReplayer(simdb::DatabaseManager* db_mgr, const std::vector<TreeNode*>& roots) :
            chkpt_loader_(db_mgr, roots)
        {
            // Load the initial simulation state
            constexpr size_t first_head_chkpt_id = 0;
            chkpt_loader_.loadCheckpoint(first_head_chkpt_id);
        }
 
        uint64_t getNumCheckpoints() const { return chkpt_loader_.getNumCheckpoints(); }
 
        bool step();
 
    private:
        DatabaseCheckpointLoader chkpt_loader_;
 
        uint64_t next_checkpoint_idx_ = 1;
    };
 
    size_t getNumCheckpoints() const;
 
    std::string stringize() const;
 
    struct ChkptWindow {
        arch_id_t start_arch_id = UINT64_MAX;
        arch_id_t end_arch_id = UINT64_MAX;
        tick_t start_tick = UINT64_MAX;
        tick_t end_tick = UINT64_MAX;
        checkpoint_ptrs checkpoints;
 
        template <typename Archive>
        void serialize(Archive & ar, const unsigned int /*version*/) {
            ar & start_arch_id;
            ar & end_arch_id;
            ar & start_tick;
            ar & end_tick;
 
            if (checkpoints.empty()) {
                // We are loading a checkpoint window from disk
                const auto num_chkpts = end_arch_id - start_arch_id + 1;
                for (size_t i = 0; i < num_chkpts; ++i) {
                    checkpoints.emplace_back(new checkpoint_type);
                    ar & *checkpoints.back();
                }
            } else {
                // We are saving a checkpoint window to disk
                for (auto& chkpt : checkpoints) {
                    ar & *chkpt;
                }
            }
        }
    };
 
    struct ChkptWindowBytes {
        arch_id_t start_arch_id = UINT64_MAX;
        arch_id_t end_arch_id = UINT64_MAX;
        tick_t start_tick = UINT64_MAX;
        tick_t end_tick = UINT64_MAX;
        std::vector<char> chkpt_bytes;
    };
 
private:
    FastCheckpointer checkpointer_;
    simdb::DatabaseManager* db_mgr_ = nullptr;
    simdb::ConcurrentQueue<ChkptWindow>* chkpt_window_head_ = nullptr;
    simdb::ConcurrentQueue<ArchIdsForTick>* tick_runs_head_ = nullptr;
    std::queue<ArchIdsForTick> tick_runs_;
    uint64_t committed_arch_id_ = 0;
    std::unique_ptr<simdb::pipeline::Flusher> pipeline_flusher_;
    std::vector<chkpt_id_t> committed_checkpoints_;
};
 
} // namespace sparta::serialization::checkpoint